[人工智能]Pytorch的常用代码小技巧

第一步：设置多卡训练正确

如果服务器的0号卡被别人占用，就得在刚开始加上这几行代码，手动更改0号显卡是2,1号卡是3。
习惯放在import os之后

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2, 3"

为什么这样写：详见
怎样选卡，超级详细？参考
训练的时候

device_ids = [0, 1]  # 这里的0和1就是前面的2和3
net = torch.nn.DataParallel(net) # , device_ids=device_ids)  # 就不加那个id选择，有多少用多少。
# 优化器的写法
optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

第二步：设置随机种子

def setup_seed(seed):
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
     np.random.seed(seed)
     random.seed(seed)
     torch.backends.cudnn.deterministic = True  # 为了确定算法，保证得到一样的结果。
     torch.backends.cudnn.enabled = True  # 使用非确定性算法
     torch.backends.cudnn.benchmark = True  # 是否自动加速。
# 设置随机数种子
setup_seed(20)

第三步：起飞小指南

就是有多少显卡，用多少显卡。

if device == 'cuda':
    net = torch.nn.DataParallel(net)
    net11 = torch.nn.DataParallel(net11)
    torch.cudnn.benchmark = True