溫馨提示×
您好，登錄后才能下訂單哦！
點擊 登錄注冊 即表示同意《億速云用戶服務條款》
您好，登錄后才能下訂單哦！
本文小編為大家詳細介紹“Pytorch模型的保存/復用/遷移怎么實現”，內容詳細，步驟清晰，細節處理妥當，希望這篇“Pytorch模型的保存/復用/遷移怎么實現”文章能幫助大家解決疑惑，下面跟著小編的思路慢慢深入，一起來學習新知識吧。
# 定義模型結構 class LenNet(nn.Module): def __init__(self): super(LenNet, self).__init__() self.conv = nn.Sequential( # [batch, 1, 28, 28] nn.Conv2d(1, 8, 5, 2), # [batch, 1, 28, 28] nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # [batch, 8, 14, 14] nn.Conv2d(8, 16, 5), # [batch, 16, 10, 10] nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # [batch, 16, 5, 5] ) self.fc = nn.Sequential( nn.Flatten(), nn.Linear(16*5*5, 128), nn.ReLU(inplace=True), nn.Linear(128, 64), nn.ReLU(inplace=True), nn.Linear(64, 10) ) def forward(self, X): return self.fc(self.conv(X))
# 查看模型參數
# 網絡模型中的參數model.state_dict()是以字典形式保存（實質上是collections模塊中的OrderedDict）
model = LenNet()
print("Model's state_dict:")
for param_tensor in model.state_dict():
print(param_tensor, "\t", model.state_dict()[param_tensor].size())
# 參數名中的fc和conv前綴是根據定義nn.Sequential()時的名字所確定。
# 參數名中的數字表示每個Sequential()中網絡層所在的位置。
print(model.state_dict().keys()) # 打印鍵
print(model.state_dict().values()) # 打印值
# 優化器optimizer的參數打印類似
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
print(var_name, "\t", optimizer.state_dict()[var_name])import os # 指定保存的模型名稱時Pytorch官方建議的后綴為.pt或者.pth model_save_dir = './model_logs/' model_save_path = os.path.join(model_save_dir, 'LeNet.pt') torch.save(model.state_dict(), model_save_path) # 在訓練過程中保存某個條件下的最優模型，可以如下操作 best_model_state = deepcopy(model.state_dict()) torch.save(best_model_state, model_save_path) # 下面這種方法是錯誤的，因為best_model_state只是model.state_dict()的引用，會隨著訓練的改變而改變 best_model_state = model.state_dict() torch.save(best_model_state, model_save_path)
def inference(data_iter, device, model_save_dir):
model = LeNet() # 初始化現有模型的權重參數
model.to(device)
model_save_path = os.path.join(model_save_dir, 'LeNet.pt')
# 如果本地存在模型，則加載本地模型參數覆蓋原有模型
if os.path.exists(model_save_path):
loaded_paras = torch.load(model_save_path)
model.load_state_dict(loaded_paras)
model.eval()
with torch.no_grad(): # 開始推理
acc_sum, n = 0., 0
for x, y in data_iter:
x, y = x.to(device), y.to(device)
logits = model(x)
acc_sum += (logits.argmax(1) == y).float().sum().item()
n += len(y)
print("Accuracy in test data is : ", acc_sum / n)class MyModel:
def __init__(self,
batch_size=64,
epochs=5,
learning_rate=0.001,
model_save_dir='./MODEL'):
self.batch_size = batch_size
self.epochs = epochs
self.learning_rate = learning_rate
self.model_save_dir = model_save_dir
self.model = LeNet()
def train(self):
train_iter, test_iter = load_dataset(self.batch_size)
# 在訓練過程中只保存網絡權重，在再訓練時只載入網絡權重參數初始化網絡訓練。這里是核心部分，開始。
if not os.path.exists(self.model_save_dir):
os.makedirs(self.model_save_dir)
model_save_path = os.path.join(self.model_save_dir, 'model.pt')
if os.path.exists(model_save_path):
loaded_paras = torch.load(model_save_path)
self.model.load_state_dict(loaded_paras)
print("#### 成功載入已有模型，進行再訓練...")
# 結束
optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.model.to(device)
for epoch in range(self.epochs):
for i, (x, y) in enumerate(train_iter):
x, y = x.to(device), y.to(device)
loss, logits = self.model(x)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 100 == 0:
acc = (logits.argmax(1) == y).float().mean()
print("Epochs[{}/{}]---batch[{}/{}]---acc {:.4}---loss {:.4}".format(
epoch, self.epochs, len(train_iter), i, acc, loss.item()))
print("Epochs[{}/{}]--acc on test {:.4}".format(epoch, self.epochs,
self.evaluate(test_iter, self.model, device)))
torch.save(self.model.state_dict(), model_save_path)
@staticmethod
def evaluate(data_iter, model, device):
with torch.no_grad():
acc_sum, n = 0.0, 0
for x, y in data_iter:
x, y = x.to(device), y.to(device)
logits = model(x)
acc_sum += (logits.argmax(1) == y).float().sum().item()
n += len(y)
return acc_sum / n# 在保存參數的時候，將優化器參數、損失值等可一同保存，然后在恢復模型時連同其它參數一起恢復
model_save_path = os.path.join(model_save_dir, 'LeNet.pt')
torch.save({
'epoch': epoch,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'loss': loss,
...
}, model_save_path)
# 加載方式如下
checkpoint = torch.load(model_save_path)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']# 定義新模型NewLeNet 和LeNet區別在于新增了一個全連接層 class NewLenNet(nn.Module): def __init__(self): super(NewLenNet, self).__init__() self.conv = nn.Sequential( # [batch, 1, 28, 28] nn.Conv2d(1, 8, 5, 2), # [batch, 1, 28, 28] nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # [batch, 8, 14, 14] nn.Conv2d(8, 16, 5), # [batch, 16, 10, 10] nn.ReLU(inplace=True), nn.MaxPool2d(2, 2), # [batch, 16, 5, 5] ) self.fc = nn.Sequential( nn.Flatten(), nn.Linear(16*5*5, 128), nn.ReLU(inplace=True), nn.Linear(128, 64), # 這層以前和LeNet結構一致 可以用LeNet的參數來進行替換 nn.ReLU(inplace=True), nn.Linear(64, 32), nn.ReLU(inplace=True), nn.Linear(32, 10) ) def forward(self, X): return self.fc(self.conv(X))
# 定義替換函數 匹配兩個網絡 size相同處地方進行參數替換
def para_state_dict(model, model_save_dir):
state_dict = deepcopy(model.state_dict())
model_save_path = os.path.join(model_save_dir, 'model.pt')
if os.path.exists(model_save_path):
loaded_paras = torch.load(model_save_path)
for key in state_dict: # 在新的網絡模型中遍歷對應參數
if key in loaded_paras and state_dict[key].size() == loaded_paras[key].size():
print("成功初始化參數:", key)
state_dict[key] = loaded_paras[key]
return state_dict# 更新一下模型遷移后的訓練代碼
def train(self):
train_iter, test_iter = load_dataset(self.batch_size)
if not os.path.exists(self.model_save_dir):
os.makedirs(self.model_save_dir)
model_save_path = os.path.join(self.model_save_dir, 'model_new.pt')
old_model = os.path.join(self.model_save_dir, 'LeNet.pt')
if os.path.exists(old_model):
state_dict = para_state_dict(self.model, self.model_save_dir) # 調用遷移代碼 將LeNet的前幾層參數遷移到NewLeNet
self.model.load_state_dict(state_dict)
print("#### 成功載入已有模型，進行再訓練...")
optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.model.to(device)
for epoch in range(self.epochs):
for i, (x, y) in enumerate(train_iter):
x, y = x.to(device), y.to(device)
loss, logits = self.model(x)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 100 == 0:
acc = (logits.argmax(1) == y).float().mean()
print("Epochs[{}/{}]---batch[{}/{}]---acc {:.4}---loss {:.4}".format(
epoch, self.epochs, len(train_iter), i, acc, loss.item()))
print("Epochs[{}/{}]--acc on test {:.4}".format(epoch, self.epochs,
self.evaluate(test_iter, self.model, device)))
torch.save(self.model.state_dict(), model_save_path)# 這里更新未進行訓練的推理
def inference(data_iter, device, model_save_dir='./MODEL'):
model = NewLeNet() # 初始化現有模型的權重參數
print("初始化參數 conv.0.bias 為：", model.state_dict()['conv.0.bias'])
model.to(device)
state_dict = para_state_dict(model, model_save_dir) # 遷移模型參數
model.load_state_dict(state_dict)
model.eval()
print("載入本地模型重新初始化 conv.0.bias 為：", model.state_dict()['conv.0.bias'])
with torch.no_grad():
acc_sum, n = 0.0, 0
for x, y in data_iter:
x, y = x.to(device), y.to(device)
logits = model(x)
acc_sum += (logits.argmax(1) == y).float().sum().item()
n += len(y)
print("Accuracy in test data is :", acc_sum / n)讀到這里，這篇“Pytorch模型的保存/復用/遷移怎么實現”文章已經介紹完畢，想要掌握這篇文章的知識點還需要大家自己動手實踐使用過才能領會，如果想了解更多相關內容的文章，歡迎關注億速云行業資訊頻道。
免責聲明：本站發布的內容（圖片、視頻和文字）以原創、轉載和分享為主，文章觀點不代表本網站立場，如果涉及侵權請聯系站長郵箱：is@yisu.com進行舉報，并提供相關證據，一經查實，將立刻刪除涉嫌侵權內容。
