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的维度。