当前位置:   article > 正文

Pytorch to(device)_pytorch to device

pytorch to device

torch.device()

  1. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  2. model.to(device)

这两行代码放在读取数据之前。 

mytensor = my_tensor.to(device)

这行代码的意思是将所有最开始读取数据时的tensor变量copy一份到device所指定的GPU上去,之后的运算都在GPU上进行。

这句话需要写的次数等于需要保存GPU上的tensor变量的个数;一般情况下这些tensor变量都是最开始读数据时的tensor变量,后面衍生的变量自然也都在GPU上。

如果是多个GPU
在代码中的使用方法为:

  1. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
  2. model = Model()
  3. if torch.cuda.device_count() > 1:
  4. model = nn.DataParallel(model,device_ids=[0,1,2])
  5. model.to(device)


torch.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model.to(device)torch.device代表将torch.Tensor分配到的设备的对象。torch.device包含一个设备类型(‘cpu’或‘cuda’)和可选的设备序号。如果设备序号不存在,则为当前设备。

如:torch.Tensor用设备构建‘cuda’的结果等同于‘cuda:X’,其中X是torch.cuda.current_device()的结果。

pytorch中model=model.to(device)用法

model=model.to(device)

这代表将模型加载到指定设备上。

其中,device=torch.device("cpu")代表的使用cpu,而device=torch.device("cuda")则代表的使用GPU

当我们指定了设备之后,就需要将模型加载到相应设备中,此时需要使用model=model.to(device),将模型加载到相应的设备中。

将由GPU保存的模型加载到CPU上。

torch.load()函数中的map_location参数设置为torch.device('cpu')

  1. device = torch.device('cpu')
  2. model = TheModelClass(*args, **kwargs)
  3. model.load_state_dict(torch.load(PATH, map_location=device))

将由GPU保存的模型加载到GPU上。确保对输入的tensors调用input = input.to(device)方法。

  1. device = torch.device("cuda")
  2. model = TheModelClass(*args, **kwargs)
  3. model.load_state_dict(torch.load(PATH))
  4. model.to(device)

将由CPU保存的模型加载到GPU上。确保对输入的tensors调用input = input.to(device)方法。map_location是将模型加载到GPU上,model.to(torch.device('cuda'))是将模型参数加载为CUDA的tensor。最后保证使用.to(torch.device('cuda'))方法将需要使用的参数放入CUDA。

  1. device = torch.device("cuda")
  2. model = TheModelClass(*args, **kwargs)
  3. model.load_state_dict(torch.load(PATH, map_location="cuda:0")) # Choose whatever GPU device number you want
  4. model.to(device)
声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/2023面试高手/article/detail/716007
推荐阅读
相关标签
  

闽ICP备14008679号