tensorflow 1.14 keras多GPU训练

虽然已经2021年了，不过最近的一个工作必须用tf1.14训练生成模型，纯tf写起来有点复杂，就用tf.keras来做，这里记录一下tf1.14中tf.keras遇到的一些问题，特别是多GPU训练

from tensorflow.keras.utils import multi_gpu_model

gpu_list = '0,1,2,3,4,5'
os.environ['CUDA_VISIBLE_DEVICES']=gpu_list #指定哪几块GPU
gpu_num = len(gpu_list.split(','))
single_model = build_model()
single_model.load_weights(path, by_name=True)
model=multi_gpu_model(single_model, gpus=gpu_num)

先指定gpu序号，建立model，加载预训练weights，调用multi_gpu_model生成多gpu训练的model。

一定要先加载权重，因为调用multi_gpu_model后的model，内部变量名什么的都变了，就加载不进来了。

 

保存中间训练结果

不能直接使用ModelCheckpoint，查了一些博客，大部分人说按照下面派生个

class ParallelModelCheckpoint(ModelCheckpoint):
    def __init__(self,model,filepath, monitor='val_loss', verbose=0,
                 save_best_only=False, save_weights_only=False,
                 mode='auto', period=1):
        self.single_model = model
        super(ParallelModelCheckpoint,self).__init__(filepath, monitor, verbose,save_best_only, save_weights_only,mode, period)

    def set_model(self, model):
        super(ParallelModelCheckpoint,self).set_model(self.single_model)

checkpoint = ParallelModelCheckpoint(model, filepath='./cifar10_resnet_ckpt.h5', monitor='val_loss', verbose=1, save_best_only=True)

可是我使用时有问题，不是每个epoch保存一次，而是每个batch保存一次，而且不能使用val_loss。

经过一顿查找，在github上找到下面方法

def get_smallest_loss(dir):
    pattern = 'model.(?P<epoch>\d+)-(?P<val_loss>[0-9]*\.?[0-9]*).h5'
    p = re.compile(pattern)
    losses = [float(p.match(f).groups()[1]) for f in os.listdir(dir) if p.match(f)]
    if len(losses) == 0:
        import sys
        return sys.float_info.max
    else:
        return np.min(losses)

class MyModelCheckpoint(Callback):
    def __init__(self, model, log_dir):
        Callback.__init__(self)
        self.model_to_save = model
        self.log_dir = log_dir
    def on_epoch_end(self, epoch, logs=None):
        fmt = self.log_dir + 'model.%04d-%.4f.h5'
        smallest_loss = get_smallest_loss(self.log_dir)
        if float(logs['val_loss']) < smallest_loss:
            self.model_to_save.save(fmt % (epoch, logs['val_loss']))

model_checkpoint = MyModelCheckpoint(single_model, log_dir)

直接从Callback派生一个，哈哈，可以每个epoch结束保存val_loss最小的中间权重了。

这里可以参考keras官方文档，https://keras.io/guides/writing_your_own_callbacks/，文档中说明了如何定义各个阶段起始结束操作。

 

======================================================================

某些硬件只支持特定的tf版本实在是太痛苦了，onnx和tf2的模型无法直接移植，只能用tf1.14重新训练一遍，如果只是检测分割还好，如果有其他复杂的，或想移植一些新op，根本搞不定