版本可参考git上的keras yolo3的实现,现在我手里有两块GPU,所以设置多块GPU来加快训练速度
1. 训练代码前引入
from keras.utils import multi_gpu_model
2. 找到自己构建网络的地方,第一次构建model是下面的语句
is_tiny_version = len(anchors) == 6 # default setting
if is_tiny_version:
model = create_tiny_model(input_shape, anchors, num_classes,
freeze_body=2, load_pretrained=False,weights_path='/home/jerry/PY_project_wang/car_detect/yolo_wang_0708/log/tiny_yolo/trained_final.h5')
else:
model = create_model(input_shape, anchors, num_classes,
freeze_body=2, weights_path='../model/old_model/yolov3_weights.h5') # make sure you know what you freeze
3. 跟踪进去,找到模型的处理,我用的tiny_yolo,在model_body 后面添加多GPU模型,我这里用了块GPU
model_body = tiny_yolo_body(image_input, num_anchors//2, num_classes)
print('Create Tiny YOLOv3 model with {} anchors and {} classes.'.format(num_anchors, num_classes))# 6 anchors and 1 classes.
model_body = multi_gpu_model(model_body, gpus=2)
4.注意,不能把multi_gpu_model()函数加在第一次构建网络的地方,即本文的2处,我试过加在哪里,结果报错如下:,为什么会这样,我也不太清楚,想来是因为最后的自定义loss层引起的
tensorflow.python.framework.errors_impl.InvalidArgumentError: Can't concatenate scalars (use tf.stack instead) for 'yolo_loss_1/concat' (op: 'ConcatV2') with input shapes: [], [], [], [].