解决 ImportError: cannot import name ‘Flatten‘ from ‘torch.nn‘

问题描述

在代码中从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)

引入后重新运行代码，没有报错