torch保存模型参数、one-hot标签、读取weights等记录-20200415

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()

参考:获取Pytorch中间某一层权重或者特征

发布了11 篇原创文章 · 获赞 0 · 访问量 315

猜你喜欢

转载自blog.csdn.net/csdnxiekai/article/details/105527813