【torch】: 并行训练并且可以动态设置第一个gpu的batch size

问题

再使用torch训练模型的时候，想要使用单机多卡

解决

参照：pytorch多gpu并行训练
需要下载自定义的:data_parallel_my_v2.py

1. 网络设置

   model = SimcseModel(pretrained_model="./model_save", pooling=args.pooler, dropout=args.dropout)

   print("torch.cuda.device_count()", torch.cuda.device_count())

   if torch.cuda.device_count() > 1:
       gpu0_bsz = int(np.sqrt(args.batch_size)) 
       # 这里设置的gpu0的batch，因为很多时候需要gpu0计算其他的东西，这个消耗会比其他gpu大
       print("gpu0_bsz", gpu0_bsz)
       acc_grad = 1 # 这里没有设置累计梯度，所以设置为1 
       print("acc_grad", acc_grad)
       from data_parallel_my_v2 import BalancedDataParallel
       model = BalancedDataParallel(gpu0_bsz // acc_grad, model, dim=0) # .cuda()

   model = model.to(args.device)

2.保存模型时设置

# 多卡gpu训练，保存方式需要做修改
if torch.cuda.device_count() > 1:
    model.module.bert.save_pretrained(save_path)
else:
    model.bert.save_pretrained(save_path)