首页 > 脚本专栏 > python > ImportError:cannot import name ‘Flatten‘ from ‘torch.nn‘

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

2023-06-15 09:09:58 作者：全幼儿园最聪明

这篇文章主要介绍了解决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)

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

以上为个人经验，希望能给大家一个参考，也希望大家多多支持脚本之家。