我正在使用 PyTorch 來訓練深度學習模型。我想知道我是否可以單獨保存模型重量。例如：
class my_model(nn.Module):
def __init__(self):
super(my_model, self).__init__()
self.bert = transformers.AutoModel.from_pretrained(BERT_PATH)
self.out = nn.Linear(768,1)
def forward(self, ids, mask, token_type):
x = self.bert(ids, mask, token_type)[1]
x = self.out(x)
return x
我有 BERT 模型作為基礎模型，頂部還有一個額外的線性層。訓練完這個模型后，我可以分別保存BERT模型和這個線性層的權重嗎？
uj5u.com熱心網友回復：
除了上一個答案，您可以創建兩個單獨的 nn.module 類。一個用于 BERT 模型，另一個用于線性層：
class bert_model(nn.Module):
def __init__(self):
super(bert_model, self).__init__()
self.bert = transformers.AutoModel.from_pretrained(BERT_PATH)
def forward(self, ids, mask, token_type):
x = self.bert(ids, mask, token_type)[1]
return x
class linear_layer(nn.Module):
def __init__(self):
super(linear_layer, self).__init__()
self.out = nn.Linear(768,1)
def forward(self, x):
x = self.out(x)
return x
然后您可以使用以下命令分別保存模型的兩部分：
bert_model = bert_model()
linear_layer = linear_layer()
#train
torch.save(bert_model.state_dict(), PATH)
torch.save(linear_layer.state_dict(), PATH)
uj5u.com熱心網友回復：
你可以：
model = my_model()
# train ...
torch.save({'bert': model.bert.state_dict(), 'out': model.out.state_dict()}, 'checkpoint.pth')
轉載請註明出處，本文鏈接：https://www.uj5u.com/qiye/381666.html
標籤：Python 机器学习 深度学习 火炬 bert语言模型
上一篇：如何在一系列日期上訓練LSTM？
下一篇：當我訓練DecisionTreeClassifier時，“輸入包含NaN、無窮大或對于dtype('float32')來說太大的值”