Tensorflow 2 单GPU同时训练多个模型

Tensorflow 2 单GPU同时训练多个模型

问题

有时我们需要对多个模型进行性能对比。若一次只训练一个模型，我们需要时刻关注训练进度，非常耗费精力。同时进行多个模型的训练能够降低人力成本。

代码

这里对三个网络进行图像的二值分割训练，它们分别是Unet, Linknet, FPN。利用for循环对整训练集进行遍历。train函数是这段代码的核心，每调用一次train就进行一次迭代。with里面是前向传播形成计算图、xxx_tape.gradient是通过反向自动微分算法求网络的梯度、xxx_optimizer.apply_gradients是优化器更新网络参数。

import tensorflow as tf
from segmentation_models import Unet,Linknet,FPN
from segmentation_models import set_framework
set_framework('tf.keras')

Unet = Unet('resnet34',input_shape=(128, 128, 3), classes=1, activation='sigmoid', encoder_weights=None, encoder_freeze=False)
Linknet = Linknet('resnet34',input_shape=(128, 128, 3), classes=1, activation='sigmoid', encoder_weights=None, encoder_freeze=False)
FPN = FPN('resnet34',input_shape=(128, 128, 3), classes=1, activation='sigmoid'