Chainerで2つの特徴マップ間の全チャネルに対する行列の内積を効率よく計算したい
chainerで画像認識をするためCNNを実装している際、ある2つの中間層の出力f1(サイズ(b1,c1,h1,w1))とf2(サイズ(b2,c2,h2,w2))に対して、同じバッチ内ですべてのc1,c2の組み合わせに対してその行列の要素積をとり和を取る、つまり行列の内積を計算するプログラムを実装し、動かしたところ、GPUメモリ上限を上回ってしまう以下のエラーが発生してしまいました。
cupy.cuda.memory.OutOfMemoryError: out of memory to allocate 2684354560 bytes (total13247891968 bytes)
よりメモリ効率の良いプログラムに改良したいのですが、よろしければアドバイスをいただけないでしょうか。以下に自分が実装したコードを載せます。エラーが出た部分は下から4行目の計算です。
環境はchainerがversion4.1、pythonが3.6、GPUにはTITAN Xを使用しています。
初質問で勝手がわからない部分もありますがよろしくお願いします。
def make_FSPmatrix(fmap1,fmap2):
b1,c1,h1,w1 = fmap1.shape
b2,c2,h2,w2 = fmap2.shape
if h1 > h2:
h,w = h2,w2
fmap1 = F.max_pooling_2d(fmap1,int(h1/h2))
elif h2 > h1:
h,w = h1,w1
fmap2 = F.max_pooling_2d(fmap2,int(h2/h1))
else:
h,w, = h1,w1
fm1_channel = F.transpose(fmap1,axes=(1,0,2,3))
fm2_channel = F.transpose(fmap2,axes=(1,0,2,3))
fsp_list = []
#内積処理
for i in range(len(fm2_channel)):
fsp_list.append(fm1_channel*F.tile(fm2_channel[i][np.newaxis,:,:],(c1,1,1,1)))
fsp_tuple = tuple(fsp_list)
fsp = F.transpose(F.sum(F.concat(fsp_tuple,axis=0),axis=(2,3)))/(h*w)
print(fsp.shape)
return fsp