开发者

解决torch.to(device)是否赋值的坑

开发者 https://www.devze.com 2024-08-10 09:47 出处:网络 作者: 不会卷积
目录torch.to(device)是否赋值的坑有一些坑需要注意,本文用RNN模型实例pytorch中model=model.to(device)用法将由GPU保存的模型加载到CPU上将由CPU保存的模型加载到GPU上总结torch.to(device)是否赋值的坑
目录
  • torch.to(device)是否赋值的坑
    • 有一些坑需要注意,本文用RNN模型实例
  • pytorch中model=model.to(device)用法
    • 将由GPU保存的模型加载到CPU上
    • 将由CPU保存的模型加载到GPU上
  • 总结

    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)

    对于模型,UKtLOJMzks不需要进行赋值:

     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_si编程客栈ze, n_class]
            return model

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

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

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

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

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

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

    device = torch.device('cpu')
    model = TheModelClass(*args, **kwargs)
    model.load_state_dict(torch.load(PATH, mhttp://www.devze.comap_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上

    确保对输入的tUKtLOJMzksensors调用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)

    总结

    以上为个人经验,希望能给大家一个参考,也希望大家多多支持编程客栈(www.devze.com)。

    0

    精彩评论

    暂无评论...
    验证码 换一张
    取 消

    关注公众号