广告位联系
返回顶部
分享到

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

python 来源:互联网 作者:佚名 发布时间:2024-06-27 21:37:29 人浏览
摘要

torch.to(device)是否赋值的坑 在我们用GPU跑程序时,需要在程序中把变量和模型放到GPU里面。 有一些坑需要注意,本文用RNN模型实例 首先,定义device 1 device = torch.device(cuda:0 if torch.cuda.is_availabl

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

在我们用GPU跑程序时,需要在程序中把变量和模型放到GPU里面。

有一些坑需要注意,本文用RNN模型实例

首先,定义device

1

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

对于变量,需要进行赋值操作才能真正转到GPU上:

1

all_input_batch=all_input_batch.to(device)

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

1

2

model = TextRNN()

model.to(device)

对模型进行to(device),还有一种方法,就是在定义模型的时候全部对模型网络参数to(device),这样就可以不需要model.to(device)这句话。

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

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 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参数设置为torch.device('cpu')

1

2

3

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)方法。

1

2

3

4

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。

1

2

3

4

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)


版权声明 : 本文内容来源于互联网或用户自行发布贡献,该文观点仅代表原作者本人。本站仅提供信息存储空间服务和不拥有所有权,不承担相关法律责任。如发现本站有涉嫌抄袭侵权, 违法违规的内容, 请发送邮件至2530232025#qq.cn(#换@)举报,一经查实,本站将立刻删除。
原文链接 :
    Tag :
相关文章
  • 本站所有内容来源于互联网或用户自行发布,本站仅提供信息存储空间服务,不拥有版权,不承担法律责任。如有侵犯您的权益,请您联系站长处理!
  • Copyright © 2017-2022 F11.CN All Rights Reserved. F11站长开发者网 版权所有 | 苏ICP备2022031554号-1 | 51LA统计