GPU一机多卡使用

我只简单写写我用的,慢慢积累.我也不可能把这块的东西全弄懂.我只给我自己看.写的很乱,sorry了.
使用多显卡的时候会自动均分.比如使用两张卡,一张是7g,一张差不多是6g.反正差不多.
使用deepspeed的时候真的是绝对均分,而且每张卡上就是一个batchsize 的大小.

#单GPU
device = torch.device("cuda:0")
model.to(device)
mytensor = my_tensor.to(device)
torch.save()
#多GPU
if torch.cuda.device_count() > 1:
     print("Let's use", torch.cuda.device_count(), "GPUs!")
     model = torch.nn.DataParallel(model)
 #多GPU保存模型
 #万能的保存方法，如果你的预测函数不会依赖你的模型类定义。
if isinstance(model, torch.nn.DataParallel):
   torch.save(model.module.state_dict(), config.save_path)
#依赖你的模型定义，这么保存只保存了你模型的参数，模型的结构没有保存，所以尽量用上面的保存方法。
if isinstance(model, torch.nn.DataParallel):
   torch.save(model.state_dict(), config.save_path)
#模型加载
model = torch.load('path/to/model')
if isinstance(model,torch.nn.DataParallel):
		model = model.module
 
#下面就可以正常使用了
model.eval()
 

和亮哥使用的一样.其实我也没全理解,凑合用吧.

    os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"#设置可见
    device = "cuda:0"#设置主卡
    model = torch.nn.DataParallel(model, device_ids=[0,1]).cuda()#设置多卡
    train_data = train_data.to(device)#模型加载

参考文献

[1] https://blog.csdn.net/pearl8899/article/details/109567745.
[2]