时间戳

首页 > 解决方案 > model.to(device) 和 model=model.to(device) 有什么区别？

python - model.to(device) 和 model=model.to(device) 有什么区别？

问题描述

假设模型本来是存储在CPU上的，然后我想把它移到GPU0上，那么我可以这样做：

device = torch.device('cuda:0')
model = model.to(device)
# or
model.to(device)

这两条线有什么区别？

标签： pythonpytorch

解决方案

没有语义差异。nn.Module.to函数将模型移动到设备。

但要小心。

对于张量（文档）：

# tensor a is in CPU
device = torch.device('cuda:0')
b = a.to(device)
# a is still in CPU!
# b is in GPU!
# a and b are different

对于模型（文档）：

# model a is in CPU
device = torch.device('cuda:0')
b = a.to(device)
# a and b are in GPU
# a and b point to the same model

推荐阅读