pytorch训练小工具收集

1、对不同的网络层配置不同的学习率

import torch 

optimizer = torch.optim.Adam([
    dict(params=model.conv1.parameters(), weight_decay=5e-4),
    dict(params=model.conv2.parameters(), weight_decay=0)
], lr=args.lr)  # Only perform weight-decay on first convolution.

2、增加维度

>>> x = torch.tensor([1, 2, 3, 4])
>>> torch.unsqueeze(x, 0)
tensor([[ 1,  2,  3,  4]])
>>> torch.unsqueeze(x, 1)
 tensor([[ 1],
        [ 2],
        [ 3],
        [ 4]])

3、可视化tensorboardX

pip install tensorboard
pip install tensorboardX