0.device的定义
就是模型和参数在CPU和GPU上的一个切换,有些地方要注意的
无非就是'cpu'和'cuda'
1.save:gpu, load:cpu
使用map_location
参数
# Specify a path to save to
PATH = "model.pt"
# Save
torch.save(net.state_dict(), PATH)
# Load
device = torch.device('cpu')
model = Net()
model.load_state_dict(torch.load(PATH, map_location=device))
2.save:gpu, load:gpu
如果有多个cuda,要指定显卡
# Save
torch.save(net.state_dict(), PATH)
# Load
device = torch.device("cuda")
model = Net()
model.load_state_dict(torch.load(PATH))
model.to(device)
注意:调用my_tensor.to(device)
只会 returns a new copy of my_tensor
on GPU (大概意思是:只会返回浅拷贝?共用内存的那种 I guess)
因此,应该重写Tensor , 代码如下:
my_tensor = my_tensor.to(torch.device('cuda')).
3. save:cpu, load:gpu
- set the
map_location
argument in thetorch.load()
function tocuda:device_id
. - call
model.to(torch.device('cuda'))
to convert the model’s parameter tensors to CUDA tensors. - use the
.to(torch.device('cuda'))
function on all model inputs to prepare the data for the CUDA optimized model.
# Save
torch.save(net.state_dict(), PATH)
# Load
device = torch.device("cuda")
model = Net()
# Choose whatever GPU device number you want
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))
# Make sure to call input = input.to(device) on any input tensors that you feed to the model
model.to(device)
4.saving torch.nn.DataParallel
Models
torch.nn.DataParalle
l is a model wrapper that enables parallel GPU utilization
(使得GPU能参与并行运算)
To save a DataParallel model generically, save the model.module.state_dict()
. This way, you have the flexibility to load the model any way you want to any device you want.
实现代码如下:
# Save
torch.save(net.module.state_dict(), PATH)
# Load to whatever device you want