前提
現在、pytorchを用いてマリオのコースをクリアするDQNのソースコードを作成し、2万回ほど学習させた時点で約3割クリアするようになっています。ここから更にクリア率を上げていきたいのですが、中断したモデルから再開しようとしても初めからになってしまっている状態です。
ですので、
- モデルを保存する際に、保存しなければならないもの
- 学習済みモデルのロードの方法
について教えていただきたいです。
なお、100epoch(episode)ごとにモデルを保存しています。
追記
モデル全体を保存するように変更したところ、学習済みモデルから再開することに成功しました。
しかし、この方法は一般的には非推奨となっており、またなぜこれで成功したのか分かっていない状態です。
ですので、成功した要因等についても教えていただきたいです。
以下のソースコードに該当部分を追記しておきます。
環境
- windows10
- python 3.10.5
- pytorch 1.12.1
実現したいこと
pytorchにおいて、中断したモデルから再度学習できるようにしたい。
セーブ・ロードをしている部分のソースコード
def save_model(self,epoch):
torch.save({
'epoch': epoch,
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict(),
'loss': self.loss},
"my_model_training_state.pt")
if os.path.exists("my_model_training_state.pt"):
print("Succeesfully load")
PATH = "my_model_training_state.pt"
checkpoint = torch.load(PATH)
self.model.load_state_dict(checkpoint['model_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
self.loss = checkpoint['loss']
self.epoch = checkpoint['epoch']
追記
保存
torch.save(model, 'save/to/path/model.pt')
読み込み
model = torch.load('load/from/path/model.pt')
補足情報（FW/ツールのバージョンなど）
以下に該当ソースコードを載せておきます。
import sys
import os
from collections import namedtuple
import numpy as np
#--------ニューラルネットワークの実装--------
import random
import copy
from copy import deepcopy
import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
#------学習に使うハイパーパラメータ---------
GAMMA = 0.9 # 時間割引率ガンマ
Transition = namedtuple('Transicion', ('state', 'action', 'next_state', 'reward'))
BATCH_SIZE = 16 # バッチサイズ
CAPACITY = 1000000 # Replay Memoryに保存するデータの最大量
EPSILON_LAST = 0.00025 #εの最低値
class ReplayMemory:
def __init__(self, CAPACITY):
self.capacity = CAPACITY #メモリの最大値
self.memory = [] # 経験を保存するリスト
self.index = 0 # 保存するindexを表す変数
def push(self, state, action, state_next, reward):
'''trasicion = (state, action, state_next, reward)をReplayMemoryに保存する'''
if len(self.memory) < self.capacity:
self.memory.append(None) # メモリが満タンじゃないときは足す
# 各引数をnamuedtupleでまとめ、memoryに保存する
self.memory[self.index] = Transition(state, action, state_next, reward)
self.index = (self.index + 1) % self.capacity # 保存するindexを1つずらす
def sample(self, batch_size):
'''batch_sizeだけ、ランダムに取り出す'''
return random.sample(self.memory, batch_size)
def __len__(self):
'''関数lenに対して、現在のmemoryの長さを返す'''
return len(self.memory)
#-----------------------------------------------------------------------------------------
class Net(nn.Module):
def __init__(self, num_states, num_actions):
super(Net, self).__init__()
self.fc1 = nn.Linear(num_states, 32)
self.fc2 = nn.Linear(32, 32) #中間層
self.fc3 = nn.Linear(32, num_actions)
def forward(self, x):
h1 = F.relu(self.fc1(x)) # 活性化関数にはReLu
h2 = F.relu(self.fc2(h1))
output = self.fc3(h2)
return output
#------------------------------------------------------------------------------------
#----行動の選択、ネットワークの更新------------------
class Brain:
def __init__(self, num_states, num_actions):
self.num_actions = num_actions # 行動の数を取得
self.action = np.zeros(5,int)
self.action_list = [
[0, 0, 0, 0, 0],#停止
[1, 0, 0, 0, 0],#左に歩く
[0, 0, 1, 0, 0],#しゃがみ
[0, 1, 0, 0, 0],#右に歩く
[0, 0, 0, 1, 0],#その場ジャンプ
[0, 1, 0, 0, 1],#右ダッシュ
[0, 1, 0, 1, 0],#右ジャンプ
[0, 1, 0, 1, 1],#右ダッシュジャンプ
[1, 0, 0, 0, 1],#左ダッシュ
[1, 0, 0, 1, 0],#左ジャンプ
[1, 0, 0, 1, 1],#左ダッシュジャンプ
]
# 経験を保存するメモリオブジェクトを生成
self.memory = ReplayMemory(CAPACITY)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
self.model = Net(num_states, num_actions)
#self.model = self.model.to(device)
self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
self.loss = None
self.epoch = 0
if os.path.exists("my_model_training_state.pt"):
print("Succeesfully load")
PATH = "my_model_training_state.pt"
checkpoint = torch.load(PATH)
self.model.load_state_dict(checkpoint['model_state_dict'])
self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
self.loss = checkpoint['loss']
self.epoch = checkpoint['epoch']
# target_net
self.target_net = copy.deepcopy(self.model)
self.target_net.load_state_dict(self.model.state_dict())
def replay(self,episode):
'''Experience Replayでネットワークの結合パラメータを出力'''
# メモリサイズがミニバッチより小さい間は何もしない
if len(self.memory) < BATCH_SIZE:
return self.loss
# メモリからミニバッチ分のデータを取り出す
transitions = self.memory.sample(BATCH_SIZE)
# 各変数をミニバッチに対応する形に変形
batch = Transition(*zip(*transitions))
# 各変数の要素をミニバッチに対応する形に変形する
state_batch = torch.cat(batch.state)
action_batch = torch.cat(batch.action)
reward_batch = torch.cat(batch.reward)
non_final_next_states = torch.cat([s for s in batch.next_state if s is not None])
# 教師信号となるQ(s_t, a_t)値を求める
self.model.eval()
# ネットワークが出力したQ(s_t, a_t)を求める
state_action_values = self.model(state_batch).gather(1, action_batch.unsqueeze(1))
# max{Q(s_t+1, a)}値を求める。ただし、次の状態があるかに注意。
# is_Finishedがdoneになっておらず、next_stateがあるかをチェックするインデックスマスクを作成
non_final_mask = torch.BoolTensor(
tuple(map(lambda s: s is not None, batch.next_state)))
# まずは全部0にしておく
next_state_values = torch.zeros(BATCH_SIZE)
# 次の状態があるindexの最大Q値を求める
self.target_net.eval()
next_state_values[non_final_mask] = self.target_net(non_final_next_states).max(1)[0].detach()
# 3.4 教師となるQ(s_t, a_t)を求める
expected_state_action_values = (next_state_values * GAMMA) + reward_batch
# ネットワークを訓練モードに切り替える
self.model.train()
# 損失関数を計算する (smooth_l1_lossはHuberloss)
self.loss = F.smooth_l1_loss(state_action_values, expected_state_action_values.unsqueeze(1))
# 結合パラメータを更新する
self.optimizer.zero_grad() # 勾配をリセット
self.loss.backward() # バックプロパゲーションを計算
self.optimizer.step() # 結合パラメータを更新
return self.loss
def save_model(self,epoch):
torch.save({
'epoch': epoch+self.epoch,
'model_state_dict': self.model.state_dict(),
'optimizer_state_dict': self.optimizer.state_dict(),
'loss': self.loss},
"my_model_training_state9.pt")
def update_target_model(self):
# モデルの重みをtarget_networkにコピー
self.target_net.load_state_dict(self.model.state_dict())
def decide_action(self, state, episode):
'''現在の状態に応じて、行動を決定する'''
epsilon = 0.5 * (1 / ((episode+self.epoch)/1000 + 1))
if epsilon < EPSILON_LAST:
epsilon = EPSILON_LAST
if epsilon <= random.random():
self.model.eval()
with torch.no_grad():
self.action = self.action_list[self.model(state).max(1)[1]]
# ネットワークの出力の最大値のindexを取り出す = max(1)[1]
# .view(1, 1)は[torch.LongTensor of size 1] を size 1x1 に変換する
else:
action_next = random.randint(0,len(self.action_list)-1)
self.action = self.action_list[action_next]
return self.action,epsilon,self.epoch
def brain_predict(self, state):
self.model.eval() # ネットワークを推論モードに切り替える
with torch.no_grad():
self.action = self.action_list[self.model(state).max(1)[1]]
return self.action
Q&A
