在使用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就可以正常使用了