2020-01-20 一些pytorch技巧

tensor的合并

cat会在原有维度上合并两个tensor
torch.cat([tensor1, tensor2], dim=0)
stack会增添一个新的维度
torch.stack([tensor1, tensor2], dim=0)

tensor的点乘

直接用*就可以了

GPU是否可用

torch.cuda.is_available()

有几块GPU

torch.cuda.device_count()

第0块GPU的名字

torch.cuda.get_device_name(0)

使用指定GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

使用GPU加速

# 将模型与损失函数转移至GPU
model =   Model()
loss_fun = torch.nn.CrossEntropyLoss()
model = model.cuda()
loss_f = loss_fun.cuda()
# 输入的数据转移至GPU
x,y = x.cuda(),y.cuda()
# 最后结果转移回CPU
loss = loss.cpu()
# 注意优化器不需要进行这一步

打印模型的参数名与参数值

# 打印参数名
for x in model.state_dict():
    print(x)
# 打印参数名与参数值
for x in model.named_parameters():
    print(x)

为什么训练模型时显存占用没问题，但在测试模型时OOM？

在测试模型的时候一定要加上with torch.no_grad(): # test_code否则被测试的模型所占用的显存会越来越大最后出现OOM

torch.manual_seed()

nn.DataParallel(model, device_ids=device_ids)

torch.nn.utils.clip_grad_norm(parameters, max_norm, norm_type=2)

tensorboardX文档

https://tensorboard-pytorch.readthedocs.io/en/latest/