进入容器:
docker exec -it 80a3cca3b305 /bin/bash
查看gpu使用情况:
nvidia-smi
查看torch和gpu配好没有:
import torch if torch.cuda.is_available(): print("gpu cuda is available!") torch.cuda.manual_seed(1000) else: print("cuda is not available! cpu is available!") torch.manual_seed(1000)
torch.cuda.manual_seed(整数)是用来给卡配好随机数种子保证一致的。
现在开始在代码里加入用来使用的gpu代码吧~
有几个分叉需要注意:要设置的是没有gpu就用cpu;有多个gpu并且可用的话就用多个gpu,否则用单gpu:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")#如果gpu没有的话就是cpu model.to(device)#然后把gpu放模型上,此时pytorch默认用1个gpu,如果要多个的话继续往下设置 n_gpu = torch.cuda.device_count()#获得当前机器有多少个gpu,它会返回一个数字,就是表示卡的数量,当然你也可以自己设置 if n_gpu > 1:#如果发现有多个gpu,那么就调用多个gpu并行计算 model = torch.nn.DataParallel(model) #当然你也可以设置一下随机数种子,如果采用多个gpu if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)#多gpu torch.cuda.manual_seed(seed)#单gpu