问题描述
在代码中从torch.nn
引入Flatten
报错:
解决方案
进入python
交互式环境
输入以下代码,查看pytorch
版本:
import torch
print(torch.__version__)
可以看到pytorch
版本是1.1.0
版本太低不包含Flatten
属性
因此,一个办法就是升级pytorch
版本,但我们知道下载安装pytorch
并不容易,若没有镜像下载缓慢且需要考虑显卡驱动、conda
、torch
三者的版本关系,升级可能导致版本不匹配conda
无法使用gpu
等问题。
另一个简单的办法是自己写Flatten
类,这里我需要的功能是仅是扁平化,参考代码如下:
class Flatten(nn.Module):
def forward(self, input):
return input.view(input.size(0), -1)
引入后重新运行代码,没有报错