AttributeError: Can‘t get attribute ‘Net‘ on module ‘__main__‘

在使用pytorch深度学习框架的时候,我们加载预先训练好的完整pkl模型时,如果报错:

AttributeError: Can't get attribute 'Net' on module '__main__'

此时我们应该在主函数内加载一个当时训练时的类,用于初始化神经网络,如下:

# 加入:初始化
class Net(nn.Module):
    def __init__(self, hidden_layers=64):
        super(Net, self).__init__()
        # hidden_layers = 128
        self.fc1    =   nn.Linear(8, hidden_layers)
        self.fc2    =   nn.Linear(hidden_layers, hidden_layers)
        self.fc_mu  =   nn.Linear(hidden_layers, 1)
    def forward(self, x):
        x   =   F.relu(self.fc1(x))
        x   =   F.relu(self.fc2(x))
        mu  =   5*torch.sigmoid(self.fc_mu(x))
        return mu
# 加入:实例化
r_net = Net(hidden_layers=64)
# 加载预先训练模型的参数
r_net = torch.load('./model.pkl')      

此时r_net就可以正常使用了

猜你喜欢

转载自blog.csdn.net/ao1886/article/details/112781534
今日推荐