MaskRCNNの学習で複数のGPUを用いるときの対処 備忘録

MaskRCNNでconfig.pyのGPU_COUNT>1にすると，バグるのでその対処方法の自分用備忘録．

バージョン・環境
OS : windows 10
GPU : GeForce GTX 1070
CUDA : 10.0 CUDA Toolkit 10.0 Archive | NVIDIA Developer
cuDNN : 7.4.1 cuDNN Archive | NVIDIA Developer

TensorFlow_version(1.14.0)
python version(3.6.5)
keras version(2.2.4)
numpy version(1.16.4)
cv2 version(4.5.1)

問題

Mask_RCNN/mrcnn/config.py
のGPU_COUNT=1→GPU_COUNT=x(x>1)とすると，

RuntimeError: It looks like you are subclassing Model and you forgot to call super(YourClass, self).__init__(). Always start with this line.

に出会える．
結論，mrcnn/parallel_model.pyのclass ParallelModel(KM.Model)を下記のように書き換える．

class ParallelModel(KM.Model):
    def __init__(self, keras_model, gpu_count):
        super(ParallelModel, self).__init__()
        self.inner_model = keras_model
        self.gpu_count = gpu_count
        merged_outputs = self.make_parallel()
        super(ParallelModel, self).__init__(inputs=self.inner_model.inputs,
                                            outputs=merged_outputs)

実際のところ，関数定義の下にsuper(ParallelModel, self).init()を挿入しているだけ．
お次は学習器のコードに戻り，学習コード前にTensorFlowが使用するGPU領域の制限を解除するコードを付け加える．

#GPUの制限解除コード
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
 
config = tf.ConfigProto()
config.allow_soft_placement=True
session = tf.Session(config=config)
KTF.set_session(session)


#学習コード
model.load_weights(COCO_MODEL_PATH, by_name=True,
                       exclude=["mrcnn_class_logits", "mrcnn_bbox_fc",
                                "mrcnn_bbox", "mrcnn_mask"])

model.train(dataset_train, dataset_val,
                learning_rate=0.01,
                epochs=1,
                layers='heads')

以上．