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：保存序列化的对象到