Pytorchを用いたFGSM実行時にmemory out of errorとなる原因

Question

#実現したいこと
Pytorchを用いたマルウェア検知のためのDeep learningフレームワークに対するFGSMを実装しようとしています．
FGSMのスクリプト実行時に，CUDA out of memoryのRuntimeErrorが発生し，最後まで実行することができません．
以下がFGSMを行なっているスクリプトです．

def fgsm_attack(self, benign_block):
        #追加するバイト列を読み込み
        append_bytez = self.benign_bytez[start_pos:end_pos]
        append_size = len(append_bytez)
        with torch.no_grad():
        #embdされた状態から元のバイト列に直すために，appendするバイト列を保存しておく
            origin_bytez = \
                torch.from_numpy(np.frombuffer(self.malware_bytez + append_bytez,\
                                               dtype=np.uint8)[np.newaxis,:] ).float()
            origin_bytez = origin_bytez.to(DEVICE)
        while True:#lossが一定以下か，ループ上限になったらbreak
            #勾配の初期化
            self.opt.zero_grad()
            #bytesの下処理．bytes -> Tensor & to -> GPU
            inp_bytez = \
                torch.from_numpy(np.frombuffer(self.malware_bytez + append_bytez,\
                                               dtype=np.uint8)[np.newaxis,:]).float()
            inp_bytez = inp_bytez.to(DEVICE)
            #gradient計算のための処理．入力byte列をembeddding
            inp_adv = inp_bytez.requires_grad_()
            inp_adv.requires_grad = True
            embd_bytez = self.model.embd_bytez(inp_adv.long())
            embd_bytez.requires_grad = True
            #modelに通して結果を出力
            out = self.model(embd_bytez)
            result = F.softmax(out, dim=1)
            #目標は良性なのでlabel=0．またlossを計算してbackprop + opimize
            label = torch.from_numpy(np.array([0]))
            label = label.to(DEVICE)
            loss = self.criterion(result, label)
            loss.backward()
            self.opt.step()
            with torch.no_grad():
                #入力バイト列全体のgradientを取得して，append部分のembd行列を取り出す
                grad = embd_bytez.grad
                grad_sign = grad.detach().sign()[0][-append_size:]
                #-> torch.size([512, 8])
            if not (grad_sign.to('cpu').numpy() == 0).all(): #all_gradient != 0の場合のみ
                #appendしているbyte列の変換，bytes -> Tensor -> embedding
                append_bytez = self.model.embd_bytez(append_bytez)
                #上のgradientを利用して，appendした部分の変更を行う
                append_bytez = (append_bytez - self.eps * grad_sign)
                #全体のembedの方でappendした部分の変更を行う
                embd_bytez[0][-append_size:] = append_bytez #update perturbation
            #embdしたバイト列を元に戻す
            recont_bytez = self.reconstruction_to_byte(origin_bytez, embd_bytez, \
                                                       append_size)
            #recont_bytezの要素をnumpy.float32 -> bytesに直す
            recont_bytez = (recont_bytez[0]).astype(np.uint8)
            recont_bytez = recont_bytez.tobytes()
            #bytes型に直下append_bytesを格納し直す
            append_bytez = recont_bytez[-append_size:]
            loop_count -= 1
            #lossが一定値以上であれば，fgsmを用いて再度更新．以下なら成でbreak
            if loop_count < 0 or loss.detach() < self.beta:
                break
        #返すのは最初のバイト列にfgsmで変更を加えたバイト列を追加したもの
        return_bytez = self.malware_bytez + append_bytez
        #メモリ不足対策に変数を消しておく
        del origin_bytez, inp_bytez, embd_bytez, loss
        return return_bytez

#試したこと
memory_profilerを用いてメモリ使用量を確認したところ，loss.backward()で大幅に増加することがあるのを確認しました．そこで変数lossに勾配情報がひたすら残り続けているのではないかと考え，最後にdelを行なったのですが，効果はありませんでした．
またbackpropに関係ない場所では，with torch.no_grad()を使用してみたのですが，これも効果がありませんでした．

そしてself.reconttuction_to_bytez()でもメモリ使用量が増えているのを確認しました．
関数の中で以下のコードを実行しているのが原因なのかな，と考えたのですが，正直そこまで問題のあるようには思えません．

 out = original_bytez.to('cpu').detach().clone().numpy()
 embd = embd_bytez.to('cpu').detach().clone().numpy()

Comments

[tiitoi]

バッチサイズを out of memory にならないまで、少なくしてはどうでしょうか？

[退会済みユーザー]

それが一回の実行で一つの入力データしか扱っていないので，バッチサイズを少なくするってことができないんですよね…

Answer 1

with torch.no_grad()の下にある，if not の部分から下の方にあるif文以下までを，with torch.no_grad()のスコープに含めると，メモリが溢れることなく最後まで動きました