基本语句

  • 模型放到GPU上运行 model.gpu() , 默认只使用一个GPU
  • 张量放到GPU上 mytensor = my_tensor.gpu()
  • 多个GPU调用： model = nn.DataParallel(model) 
  1. model = Model(input_size, output_size)
  2. if torch.cuda.device_count() > 1:
  3.   print("Let's use", torch.cuda.device_count(), "GPUs!")
  4.   # dim = 0 [30, xxx] -> [10, ...], [10, ...], [10, ...] on 3 GPUs
  5.   model = nn.DataParallel(model)
  7. if torch.cuda.is_available():
  8.    model.cuda()
  10. for data in rand_loader:
  11.     if torch.cuda.is_available():
  12.         input_var = Variable(data.cuda())
  13.     else:
  14.         input_var = Variable(data)
  16.     output = model(input_var)
  17.     print("Outside: input size", input_var.size(),
  18.           "output_size", output.size())