前回の日記で、RL policy networkの勾配を求める際に、対数尤度の偏微分に報酬に応じた重み(勝敗の報酬から状態価値を引いた値)を掛ける計算の実装が、Chainerでは難しいということを書いた。
Chainerでは損失関数のbackwardを行うと、ミニバッチで1つの勾配が計算されるため、ミニバッチの要素(局面)単位の勾配にを掛ける計算が実装できない。
そこで、ミニバッチの要素1件ずつ、順伝播と逆伝播を行い、計算された勾配を保存しておき、ミニバッチのすべての要素の勾配を計算できたら、保存しておいた勾配に、報酬に応じた重みを掛け合わせて、それらの平均をとることでミニバッチの勾配とする。
最後に、モデルのupdateを行い、パラメータを更新する。
この方法で、報酬に応じた勾配が実装できる。
しかし、ミニバッチの単位で順伝播、逆伝播の計算ができないため、1件ずつ計算すると時間がかかってしまう。
この部分を効率化するにはChainerの内部で行っている勾配計算に手を入れるか、TensorFlowなどでより低レベルな実装を行う必要がある。
また、勾配の計算はGPUからCPUに値を移して計算する必要があるため、メモリ転送で時間のロスが発生する。
実行時間を気にしなければ、以下のように記述することで実装できる。
下記の例は、3層パーセプトロンで、ミニバッチの要素ごとに勾配に重みを掛けている。
import numpy as np from chainer import cuda, Chain, Function, Variable from chainer import optimizers import chainer.functions as F import chainer.links as L W1 = [[ 1.21082544, -0.42751756], [ 1.35623264, -0.1971387 ], [-0.77883673, 0.28367677]] W2 = [[ 0.08621028, -0.19540818, 0.78203094], [ 0.30133799, 1.3698988 , -0.01031571]] class MyChain(Chain): def __init__(self): super(MyChain, self).__init__( l1=L.Linear(2, 3, initialW=np.array(W1)), l2=L.Linear(3, 2, initialW=np.array(W2)), ) def __call__(self, x): h = self.l1(x) return self.l2(h) model = MyChain() model.to_gpu() optimizer = optimizers.SGD() optimizer.setup(model) # print param data for path, param in model.namedparams(): print(path) print(param.data) print() x_data = [[1, 2], [3, 4]] t_data = [0, 1] # init grads dic_grads = {} for path, param in model.namedparams(): dic_grads[path] = [] for x_elem, t_elem in zip(x_data, t_data): x = Variable(cuda.to_gpu(np.array([x_elem], dtype=np.float32))) t = Variable(cuda.to_gpu(np.array([t_elem], dtype=np.int32))) y = model(x) model.cleargrads() loss = F.softmax_cross_entropy(y, t) loss.backward() # save grad for path, param in model.namedparams(): dic_grads[path].append(cuda.to_cpu(param.grad)) # manipulate grad z_data = [1.0, 0.5] z = np.array(z_data, dtype=np.float32) for path, param in model.namedparams(): grads = np.array(dic_grads[path], dtype=np.float32) grad = np.tensordot(z, grads, axes=1) grad /= len(z_data) param.grad = cuda.to_gpu(grad) optimizer.update() # print param data and grad for path, param in model.namedparams(): print(path) print(cuda.to_cpu(param.data)) print(cuda.to_cpu(param.grad))
zがミニバッチの要素(要素数2)ごとの重みとなっている。
z_data = [1.0, 1.0]
とすると、Chainerの勾配計算と同じになる。
上記の方法ではGPUを活かしきれず非効率であるため、Chainerの内部処理をもう少し理解して、効率化を検討する予定。
softmax_cross_entropyのbackwardに手を入れることで実現できないか調べている。