Multi-GPU Model에서 h5(hdf5)를 로드하지 못하는 문제

문제 개요

Multi-GPU Model로 학습시키고 h5(hdf5)로 저장해둔 다음, 나중에 불러올 때 layers 수가 맞지 않다고 에러가 남

이미 많이 알려진 버그

 

해결책 개요

Multi-GPU Model로 학습한 weight를 (Multi-GPU로 compile하지 않은) Base Model에 set한 다음 h5(hdf5)로 저장

 

Epoch가 모두 끝난 다음 model.save()는 대략 아래와 같이 하면 되지만...

...
parallel_model, base_model = get_compiled_model(...)

parallel_model.fit(...)

base_model.set_weights(parallel_model.get_weights())
base_model.save(filepath)
...

Keras에서 제공하는 Callback 중 ModelCheckpoint를 사용하려면 간단히 바꿀 수 있는 방법이 없다...

그래서 ModelCheckpoint를 상속받아서 위와 같이 약간 수정해줘야 한다. 아래 링크들을 참고..

 

https://github.com/keras-team/keras/blob/master/keras/callbacks.py#L275

https://github.com/keras-team/keras/blob/master/keras/callbacks.py#L633

순서대로 Callback, ModelCheckpoint 코드

 

https://github.com/keras-team/keras/issues/8123#issuecomment-348976624

ModelCheckpoint Custom rf

 

그래서 바꾼 Multi-GPU Model 용 ModelCheckpoint는 아래와 같다.

class MultiGPUModelCheckpoint(Callback):
    """
    Multi GPU 사용 시 모델 저장 후 로드가 안 되어 base_model을 통해 저장함
    """
    def __init__(self, filepath, base_model, monitor='val_loss', verbose=0,
                 save_best_only=False, save_weights_only=False,
                 mode='auto', period=1):
        super(MultiGPUModelCheckpoint, self).__init__()
        self.monitor = monitor
        self.verbose = verbose
        self.filepath = filepath
        self.save_best_only = save_best_only
        self.save_weights_only = save_weights_only
        self.period = period
        self.epochs_since_last_save = 0
        self.base_model = base_model

        if mode not in ['auto', 'min', 'max']:
            warnings.warn('ModelCheckpoint mode %s is unknown, '
                          'fallback to auto mode.' % (mode),
                          RuntimeWarning)
            mode = 'auto'

        if mode == 'min':
            self.monitor_op = np.less
            self.best = np.Inf
        elif mode == 'max':
            self.monitor_op = np.greater
            self.best = -np.Inf
        else:
            if 'acc' in self.monitor or self.monitor.startswith('fmeasure'):
                self.monitor_op = np.greater
                self.best = -np.Inf
            else:
                self.monitor_op = np.less
                self.best = np.Inf
    
    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        self.epochs_since_last_save += 1
        if self.epochs_since_last_save >= self.period:
            self.epochs_since_last_save = 0
            filepath = self.filepath.format(epoch=epoch + 1, **logs)
            if self.save_best_only:
                current = logs.get(self.monitor)
                if current is None:
                    warnings.warn('Can save best model only with %s available, '
                                  'skipping.' % (self.monitor), RuntimeWarning)
                else:
                    if self.monitor_op(current, self.best):
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s improved from %0.5f to %0.5f,'
                                  ' saving model to %s'
                                  % (epoch + 1, self.monitor, self.best,
                                     current, filepath))
                        self.best = current
                        if self.save_weights_only:
                            self.base_model.set_weights(self.model.get_weights())
                            self.base_model.save_weights(filepath, overwrite=True)
                        else:
                            self.base_model.set_weights(self.model.get_weights())
                            self.base_model.save(filepath, overwrite=True)
                    else:
                        if self.verbose > 0:
                            print('\nEpoch %05d: %s did not improve from %0.5f' %
                                  (epoch + 1, self.monitor, self.best))
            else:
                if self.verbose > 0:
                    print('\nEpoch %05d: saving model to %s' % (epoch + 1, filepath))
                if self.save_weights_only:
                    self.base_model.set_weights(self.model.get_weights())
                    self.base_model.save_weights(filepath, overwrite=True)
                else:
                    self.base_model.set_weights(self.model.get_weights())
                    self.base_model.save(filepath, overwrite=True)

hdf5 파일로 잘 저장되는 것과, 다시 로드하는 것이 되는 것을 확인했다.