pytorch常用函数理解

本博客用来记录常用的torch中的函数

torch.max

output = torch.max(x,dim=1)

• input输入的是一个tensor

• dim是max函数索引的维度0/1，0是每列的最大值，1是每行的最大值

• 返回的是两个值：一个是每一行或列最大值的tensor组，另一个是最大值所在的位置(索引）

x = torch.rand(3,5)
# tensor([[0.3168, 0.8239, 0.1350, 0.8047, 0.8741],
#         [0.8606, 0.4046, 0.3248, 0.0284, 0.8336],
#         [0.3513, 0.4505, 0.7712, 0.2254, 0.2427]])

torch.max(x,0) # 返回值包括每一列的最大值和最大值所在的索引
# torch.return_types.max(
# values=tensor([0.8606, 0.8239, 0.7712, 0.8047, 0.8741]),
# indices=tensor([1, 0, 2, 0, 0]))

torch.max(x,1) # 返回值包括每一行的最大值和最大值所在的索引
# torch.return_types.max(
# values=tensor([0.8741, 0.8606, 0.7712]),
# indices=tensor([4, 0, 2]))

torch.max(x,0)[0]  # 返回每一列的最大值
# tensor([0.8606, 0.8239, 0.7712, 0.8047, 0.8741])

torch.max(x,0)[1] # 返回每一列最大值的索引
# tensor([1, 0, 2, 0, 0])

pytorch的模型保存与加载

• torch.save：保存序列化的对象到磁盘，使用了Python的pickle进行序列化，模型、张量、所有对象的字典。
• torch.load：使用了pickle的unpacking将pickled的对象反序列化到内存中。
• torch.nn.Module.load_state_dict：使用反序列化的state_dict加载模型的参数字典。

1、保存/加载state_dict{}

# 模型保存
torch.save(model.state_dict(), PATH)

# 模型加载
model = ModelNet(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

 保存模型的推理过程的时候，只需要保存模型训练好的参数，使用torch.save()保存state_dict，能够方便模型的加载。因此推荐使用这种方式进行模型保存。

PyTorch 通用模型保存格式为 .pt 和 .pth 文件扩展名形式.

注意，load_state_dict()需要传入字典对象，因此需要先反序列化state_dict再传入load_state_dict()

 2、保存/加载整个模型

# 保存
torch.save(model, PATH)

# 加载
model = ModelNet(*args, **kwargs) # 必须预先定义过模型.
model = torch.load(PATH)
model.eval()

这种方式是最直观的语法，包含最少的代码. 其会采用 Python 的pickle 模块保存全部的模型模块.

这种方法的缺点是，保存模型的时候，序列化的数据被绑定到了特定的类和确切的目录。这是因为pickle不保存模型类本身，而是保存这个类的路径，并且在加载的时候会使用。因此，当在其他项目里使用或者重构的时候，加载模型的时候会出错。

3、保存/加载 CHECKPOINT 用于推断或恢复训练

# 保存
torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

# 加载
model = ModelNet(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# 或
model.train()

在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。保存优化器的state_dict也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时epoch数，最新的模型损失，额外的torch.nn.Embedding层等。

要保存多个组件，则将它们放到一个字典中，然后使用torch.save()序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。

加载各个组件，首先初始化模型和优化器，然后使用torch.load()加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。

4、保存多个模型到一个文件

# 保存
torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)

# 加载
modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - or -
modelA.train()
modelB.train()

当需要将多个 torch.nn.Modules 组成的模型进行保存时，如 GAN，sequence-to-sequence 模型，模型集成ensembleofmodels，可以采用与保存断点checkpoints模型相同的方法.

换句话说，将每个模型的 state_dict 以及对应的优化器保存为一个字典.

还可以保存任何其它有助于恢复训练的信息项，只需简单的添加其到字典中.