解说pytorch中的model=model.to(device)

共计 2868 个字符，预计需要花费 8 分钟才能阅读完成。

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

其中，device=torch.device(“cpu”) 代表的使用 cpu，而 device=torch.device(“cuda”) 则代表的使用 GPU。

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

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

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

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

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

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
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。

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)

补充：pytorch 中 model.to(device) 和 map_location=device 的区别

在已训练并保存在 CPU 上的 GPU 上加载模型时，加载模型时经常由于训练和保存模型时设备不同出现读取模型时出现错误，在对跨设备的模型读取时候涉及到两个参数的使用，分别是 model.to(device) 和 map_location=devicel 两个参数，简介一下两者的不同。

将 map_location 函数中的参数设置 torch.load() 为 cuda：device_id。这会将模型加载到给定的 GPU 设备。

调用 model.to(torch.device(‘cuda’)) 将模型的参数张量转换为 CUDA 张量，无论在 cpu 上训练还是 gpu 上训练，保存的模型参数都是参数张量不是 cuda 张量，因此，cpu 设备上不需要使用 torch.to(torch.device(“cpu”))。

了解了两者代表的意义，以下介绍两者的使用。

保存：

torch.save(model.state_dict(), PATH)

加载：

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

解释：

在使用 GPU 训练的 CPU 上加载模型时，请传递 torch.device(‘cpu’) 给 map_location 函数中的 torch.load() 参数，使用 map_location 参数将张量下面的存储器动态地重新映射到 CPU 设备。

保存：

torch.save(model.state_dict(), PATH)

加载：

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)# Make sure to call input = input.to(device) on any input tensors that you feed to the model

解释：

在 GPU 上训练并保存在 GPU 上的模型时，只需将初始化 model 模型转换为 CUDA 优化模型即可 model.to(torch.device(‘cuda’))。

此外，请务必.to(torch.device(‘cuda’)) 在所有模型输入上使用该 功能来准备模型的数据。

请注意，调用 my_tensor.to(device) 返回 my_tensorGPU 上的新副本。

它不会覆盖 my_tensor。

因此，请记住手动覆盖张量：my_tensor = my_tensor.to(torch.device(‘cuda’))

保存：

torch.save(model.state_dict(), PATH)

加载：

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # Choose whatever GPU device number you want
model.to(device)# Make sure to call input = input.to(device) on any input tensors that you feed to the model

解释：

在已训练并保存在 CPU 上的 GPU 上加载模型时，请将 map_location 函数中的参数设置 torch.load() 为 cuda：device_id。

这会将模型加载到给定的 GPU 设备。

接下来，请务必调用 model.to(torch.device(‘cuda’)) 将模型的参数张量转换为 CUDA 张量。

最后，确保.to(torch.device(‘cuda’)) 在所有模型输入上使用该 函数来为 CUDA 优化模型准备数据。

请注意，调用 my_tensor.to(device) 返回 my_tensorGPU 上的新副本。

它不会覆盖 my_tensor。

因此，请记住手动覆盖张量：my_tensor = my_tensor.to(torch.device(‘cuda’))

阿里云 2 核 2G 服务器 3M 带宽 61 元 1 年，有高配

腾讯云新客低至 82 元 / 年，老客户 99 元 / 年

代金券：在阿里云专用满减优惠券