解決torch.to(device)是否賦值的坑

更新時間：2024年06月27日 14:45:39 作者：不會卷積

這篇文章主要介紹了解決torch.to(device)是否賦值的坑,具有很好的參考價值,希望對大家有所幫助,如有錯誤或未考慮完全的地方,望不吝賜教

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

在我們用GPU跑程序時，需要在程序中把變量和模型放到GPU里面。

有一些坑需要注意，本文用RNN模型實例

首先，定義device

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

對于變量，需要進行賦值操作才能真正轉(zhuǎn)到GPU上：

all_input_batch=all_input_batch.to(device)

對于模型，不需要進行賦值：

 model = TextRNN()
 model.to(device)

對模型進行to(device)，還有一種方法，就是在定義模型的時候全部對模型網(wǎng)絡(luò)參數(shù)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 model

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

這代表將模型加載到指定設(shè)備上。

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

當(dāng)我們指定了設(shè)備之后，就需要將模型加載到相應(yīng)設(shè)備中，此時需要使用model=model.to(device)，將模型加載到相應(yīng)的設(shè)備中。

將由GPU保存的模型加載到CPU上

將torch.load()函數(shù)中的map_location參數(shù)設(shè)置為torch.device('cpu')

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

將由GPU保存的模型加載到GPU上。確保對輸入的tensors調(diào)用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調(diào)用input = input.to(device)方法。

map_location是將模型加載到GPU上，model.to(torch.device('cuda'))是將模型參數(shù)加載為CUDA的tensor。

最后保證使用.to(torch.device('cuda'))方法將需要使用的參數(shù)放入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)