把一个py文件中的网络应用于另一个py文件后出错

问题

训练了一个自编码器网络，想把训练好的模型用在另一个文件里，结果报了AttributeError: Can’t get attribute ‘AutoEncoder’ on <module ‘main’ from ‘D:/alicia_project/mycutpaste/train.py’>的错误

训练好的模型文件（部分）：

## ae.py
## 网络结构
class AutoEncoder(nn.Module):
    def __init__(self):
        super(AutoEncoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.Tanh(),
            nn.Linear(128, 64),
            nn.Tanh(),
            nn.Linear(64, 32),
            nn.Tanh(),
            nn.Linear(32, 16),
            nn.Tanh(),
            nn.Linear(16, 3)
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 16),
            nn.Tanh(),
            nn.Linear(16, 32),
            nn.Tanh(),
            nn.Linear(32, 64),
            nn.Tanh(),
            nn.Linear(64, 128),
            nn.Tanh(),
            nn.Linear(128, 28 * 28),
            nn.Sigmoid()

        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return encoded, decoded
## 实例化
Coder = AutoEncoder()
## 加载训练好的模型
Coder = torch.load('AutoEncoder.pkl')

当我把此网络引用在另一个文件里（部分代码）：

## train.py
import ae
## 应用模型
_, result = ae.Coder(patch_resize)

结果报错AttributeError: Can’t get attribute ‘AutoEncoder’ on <module ‘main’ from ‘D:/alicia_project/mycutpaste/train.py’>

原因

可以参考这篇博客，讲的很详细。训练好的pkl文件,再次加载的时候AttributeError: Can‘t get attribute ‘Model‘ on ＜module ‘main‘＞_pkl文件加载_ERCO123的博客-优快云博客

解决

移植网络

上网查阅了一些资料，最后有用的是，不仅需要引用模块 import ae还需要把整个网络也移植过去，train.py更改为如下：

## train.py
## 网络结构
class AutoEncoder(nn.Module):
    def __init__(self):
        super(AutoEncoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(28 * 28, 128),
            nn.Tanh(),
            nn.Linear(128, 64),
            nn.Tanh(),
            nn.Linear(64, 32),
            nn.Tanh(),
            nn.Linear(32, 16),
            nn.Tanh(),
            nn.Linear(16, 3)
        )
        self.decoder = nn.Sequential(
            nn.Linear(3, 16),
            nn.Tanh(),
            nn.Linear(16, 32),
            nn.Tanh(),
            nn.Linear(32, 64),
            nn.Tanh(),
            nn.Linear(64, 128),
            nn.Tanh(),
            nn.Linear(128, 28 * 28),
            nn.Sigmoid()

        )

    def forward(self, x):
        encoded = self.encoder(x)
        decoded = self.decoder(encoded)
        return encoded, decoded
import ae
## 应用模型
_, result = ae.Coder(patch_resize)

保存并加载权重

如果已经移植了网络，还是会报错的话，就不要直接保存模型了，参考这篇博客 AttributeError: Can’t get attribute ‘Net’ on module ‘main’_程序员阿德的博客-优快云博客，直接保存网络权重参数。