tensorflowでバッチサイズを下げて運用する方法

Question

以下のサイトから超解像のプログラムを取得してきました．
https://github.com/Saafke/EDSR_Tensorflow

そのrun.pyのupscaleFromPb関数

def upscaleFromPb(self, path):
        """
        Upscale single image by desired model. This loads a .pb file.
        """
        # Read model
        pbPath = "./models/EDSR_x{}.pb".format(self.scale)

        # Get graph
        graph = self.load_pb(pbPath)

        fullimg = cv2.imread(path, 3)
        floatimg = fullimg.astype(np.float32) - self.mean
        LR_input_ = floatimg.reshape(1, floatimg.shape[0], floatimg.shape[1], 3)

        LR_tensor = graph.get_tensor_by_name("IteratorGetNext:0")
        HR_tensor = graph.get_tensor_by_name("NHWC_output:0")

        with tf.compat.v1.Session(graph=graph) as sess:
            print("Loading pb...")
            output = sess.run(HR_tensor, feed_dict={LR_tensor: LR_input_})
            Y = output[0]
            HR_image = (Y + self.mean).clip(min=0, max=255)
            HR_image = (HR_image).astype(np.uint8)

            cv2.imwrite("./finish.jpg", HR_image)
            cv2.destroyAllWindows()

        sess.close()

のoutput = sess.run(HR_tensor, feed_dict={LR_tensor: LR_input_})部分で

tensorflow.python.framework.errors_impl.ResourceExhaustedError: 2 root error(s) found.
  (0) Resource exhausted: OOM when allocating tensor with shape[1,256,2480,3508] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
         [[node Conv2D (defined at D:\Downloads\EDSR_Tensorflow-master\EDSR_Tensorflow-master\run.py:224) ]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info. This isn't available when running in Eager mode.

が発生します．

https://qiita.com/enoughspacefor/items/1c09a27877877c56f25aの方が，同じエラーでバッチサイズを小さくするとあるのですが，どの部分がバッチサイズに当たるのかわかりません．またバッチサイズを小さくするとfor文などで，繰り返しsess.runを行わないといけないのでしょうか．

Answer 1

解説に

Extra arguments (Nr of resblocks, filters, batch, lr etc.) python main.py --help

とあります。

これは「エキストラ引数（resblockの数、フィルタ、バッチ、学習率など）はpython main.py --helpで確認できます。」の意味です。

# main.py
parser.add_argument('--batch', type=int, help='Batch size of the training', default=16)

とありますので、python main.py --batch 4のようにすれば標準の16よりも小さいバッチで動かせます。これでOOMは防げそうですね。

Comments

[takugon92]

ありがとうございます．
こちらの説明不足でした．
python main.py --batchの指定は学習をさせる際には有効なのですが，私がエラーが出ているのは，画像ファイルを学習済みデータを使って超解像を試しているときに出ています．
このとき使っているのがupscaleFromPb(self, path)なんですが，ここには--batchで指定した数値が効いて来ないみたいです．

400x400の小さい画像ならできたので，3508ｘ2480の大きい画像で試してみたらエラーが出るようになりました．

[退会済みユーザー]

400x400,600x600,800x800,1000x1000といった具合で画像サイズとメモリの使用量を散布図にしてはいかがでしょうか？恐らく指数関数的にメモリを喰うはずなので3000x2000だとOOMが起きる、は十分あると思います。
…このモデルには詳しくないですが、普通trainの時にはbachは有効ですが、Estimateのbatchは（複数枚を同時にEstimateする場合を除き）普通1で処理されます。

参考：TITAN X (12GB)でも3000x2000クラスはOOMするらしいです。
https://github.com/clovaai/cutblur/issues/4

[退会済みユーザー]

やったことはないですが、ムリにNNで解像度を上げるより、画像を分割してそれぞれアップスケーリングなんてどうなんでしょうね。継ぎ目が荒れちゃうかも知れませんが…

[takugon92]

ありがとうございます．まさに仰るとおり800x800以降でOOMが発生しました．
タイタンでも発生するのですね．2080Sなので無論発生しますね．

GPUを使わず，CPUで無理やり引っ張ると（CPU，メモリ使用率共に100％），ryzen9 3950xと2080Sの組み合わせで，3000x2000はギリギリ解像度を上げることができました．
3508ｘ2480ですと，途中でエラーも発生せず止まりました↓
https://teratail.com/questions/91121

画像を分割してそれぞれアップスケーリングなんてどうなんでしょうね。
⇒やはりこの方法しかなさそうですね．確かにつなぎ目が荒れるかもしれませんが一度やってみます．
相談に乗ってくださり，ありがとうございました．

[takugon92]

分割し，実行してみました．
つなぎ目のあれなど特になく，超解像で4倍にアップスケーリングできました．
ありがとうございました