我要评论torch.to(device)是否赋值的坑
在我们用gpu跑程序时，需要在程序中把变量和模型放到gpu里面。
有一些坑需要注意，本文用rnn模型实例
首先，定义device
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")对于变量，需要进行赋值操作才能真正转到gpu上：
all_input_batch=all_input_batch.to(device)
对于模型，不需要进行赋值：
model = textrnn() model.to(device)
对模型进行to(device)，还有一种方法，就是在定义模型的时候全部对模型网络参数to(device)，这样就可以不需要model.to(device)这句话。
class textrnn(nn.module):
def __init__(self):
super(textrnn, self).__init__()
#self.cnt = 0
self.c = nn.embedding(n_class, embedding_dim=emb_size,device=device)
self.rnn = nn.rnn(input_size=emb_size, hidden_size=n_hidden,device=device)
self.w = nn.linear(n_hidden, n_class, bias=false,device=device)
self.b = nn.parameter(torch.ones([n_class])).to(device)
def forward(self, x):
x = self.c(x)
#print(x.is_cuda)
x = x.transpose(0, 1) # x : [n_step, batch_size, embeding size]
outputs, hidden = self.rnn(x)
# outputs : [n_step, batch_size, num_directions(=1) * n_hidden]
# hidden : [num_layers(=1) * num_directions(=1), batch_size, n_hidden]
outputs = outputs[-1] # [batch_size, num_directions(=1) * n_hidden]
model = self.w(outputs) + self.b # model : [batch_size, n_class]
return modelpytorch中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')
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)总结
以上为个人经验，希望能给大家一个参考，也希望大家多多支持代码网。
发表评论