save:
torch.save(model.state_dict(), PATH)
在pytorch保存模型进行推理时,只需要保存训练过的模型的学习参数即可。一个常见的PyTorch约定是使用.pt或.pth文件扩展名保存模型。
load:
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
记住,您必须调用model.eval(),以便在运行推断之前将dropout和batch规范化层设置为评估模式。如果不这样做,将会产生不一致的推断结果。
Note:
注意,load_state_dict()函数接受一个dictionary对象,而不是保存对象的路径。这意味着您必须在将保存的state_dict传至load_state_dict()函数之前反序列化它。
参考:PyTorch之保存加载模型
one-hot:
旧版:
class_num = 10
batch_size = 4
#label = torch.LongTensor(batch_size, 1).random_() % class_num
one_hot = torch.zeros(batch_size, class_num).scatter_(1, label, 1)
新版:
torch.nn.functional.one_hot
F.one_hot(torch.arange(0, 5), num_classes=5)
参考:PyTorch——Tensor_把索引标签转换成one-hot标签表示(包括回复)
读取中间某一层权重weight和bias
parm={}
for name,parameters in G.named_parameters():
print(name,':',parameters.size())
parm[name]=parameters.cpu().detach().numpy()