pytorch常用语句

1. 检测是否存在cuda

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # PyTorch v0.4.0
model = CNN().to(device)

2. numpy与torch互转

# numpy 转为 torch
torch_data = torch.from_numpy(np_data)

# torch 转为numpy
tensor2array = torch_data.numpy()

3. 打印网络形状

需要提前安装torchsummary库，直接pip安装就是可以的。

from torchsummary import summary

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = CNN().to(device)
    summary(model, (1, 3)) #（1，3）为模型输入

4. 保存加载权重

保存权重

# 建立一个字典，保存三个参数
state = {
    ‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}
# 调用torch.save()
# dir表示保存文件的绝对路径+保存文件名，如'/home/qinying/Desktop/modelpara.pth'
torch.save(state, dir)

加载权重

checkpoint = torch.load(dir)

model.load_state_dict(checkpoint['net'])

optimizer.load_state_dict(checkpoint['optimizer'])

start_epoch = checkpoint['epoch'] + 1

5. 数据维度改变

5.1 维度压缩

使用函数torch.squeeze()对数据的维度进行压缩，去掉维数为1的的维度，比如是一行或者一列这种，一个一行三列（1,3）的数去掉第一个维数为一的维度之后就变成（3）行。squeeze(a)就是将a中所有为1的维度删掉。不为1的维度没有影响。a.squeeze(N) 就是去掉a中指定的维数为一的维度。
还有一种形式就是b=torch.squeeze(a，N) a中去掉指定的定的维数为一的维度。

5.2 维度扩充

torch.unsqueeze()函数主要是对数据维度进行扩充。给指定位置加上维数为一的维度，比如原本有个三行的数据（3），在0的位置加了一维就变成一行三列（1,3）。a.squeeze(N) 就是在a中指定位置N加上一个维数为1的维度。还有一种形式就是b=torch.squeeze(a，N) a就是在a中指定位置N加上一个维数为1的维度。