最近發現

在深度學習模型中

不能在模型建構中

先將參數轉成 cuda 格式

例如以下這樣

正常來說建構模型是以下這樣

import torch
import torch.nn as nn
import torch.optim as optim

class CustomModel(nn.Module):
    def init(self):
       super(CustomModel, self).init()
       self.linear = nn.Linear(10, 1)

        # 定義自定義參數
        self.custom_param = nn.Parameter(torch.randn(1))

    def forward(self, x):
        x = self.linear(x)
        return x + self.custom_param

# 創建模型實例
model = CustomModel().cuda() # 將模型移動到 GPU

### 訓練後...


#儲存權重 
torch.save(model.state_dict(), 'custom_model.pth')


# 創建新的模型實例
new_model = CustomModel().cuda() # 新模型也需要移動到 GPU 
# 加載儲存的權重 
new_model.load_state_dict(torch.load('custom_model.pth'))


# 確保新模型進入評估模式 
new_model.eval()

但有些時候自己沒注意到

變成以下這樣

class CustomModel(nn.Module):
    def init(self):
       super(CustomModel, self).init()
       self.linear = nn.Linear(10, 1)

        # 定義自定義參數
        self.custom_param = nn.Parameter(torch.randn(1)).cuda() # 這樣有問題

    def forward(self, x):
        x = self.linear(x)
        return x + self.custom_param

這樣也能夠訓練

也能夠儲存權重

也不會報錯

但會發現 custom_param 實際上並沒有真的寫入到權重中

所以要避免這種宣告方式喔~

 

給大家參考囉