在多GPU环境下的Keras中使用tf.image.resize_bilinear时获得InvalidArgumentError

Question

我在分段网络中使用tf.image.resize_bilinear，看来此功能在多GPU模型中不支持。以下代码显示了简化的情况：（可以直接运行）

import os
os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1'
from keras.backend.tensorflow_backend import set_session
from keras import backend as K
from keras.utils import multi_gpu_model
from keras.applications.mobilenet_v2 import preprocess_input
import tensorflow as tf
import numpy as np

config = tf.ConfigProto()
config.gpu_options.allow_growth = True  
config.allow_soft_placement = True         
sess = tf.Session(config=config)
set_session(sess)

batch = 4
num_classes = 2
size = 128
K.clear_session()

def _GetRandomImg():
    shape = (batch, size, size, 3)
    img = np.random.randint(low=0, high=256, size=shape)
    return preprocess_input(img)

def _GetRandomLabel():
    shape = (batch, size, size, num_classes)
    label = np.random.randint(low=0, high=num_classes, size=shape)
    label = np.exp(label)
    label = label/ np.sum(label, axis=-1, keepdims=True)
    return label

def DataGen():
    while True:
        x = _GetRandomImg()
        y = _GetRandomLabel()
        yield x, y

from keras.layers import Input, Conv2D, Lambda
from keras import Model

def GetModel():
    inputs = Input(shape=(size, size, 3))
    f = lambda x: tf.image.resize_bilinear(inputs, (size, size), align_corners=True)
    x = Lambda(f, output_shape=(size, size, 3))(inputs)
    outputs = Conv2D(num_classes, kernel_size=3, padding='same')(x)
    model = Model(inputs=[inputs], outputs=[outputs])
    return model

gen = DataGen()
with tf.device('/cpu:0'):
    model = GetModel() 
model = multi_gpu_model(model, gpus=2)
model.compile(loss='categorical_crossentropy', optimizer='sgd')
result = model.fit_generator(gen, epochs=2, verbose = 1, steps_per_epoch = 100)

它在单GPU环境下工作正常，但是在多GPU环境下，出现以下错误：

InvalidArgumentError: Incompatible shapes: [3,128,128,2] vs. [6,128,128,2]
     [[{{node loss/conv2d_1_loss/categorical_crossentropy/mul}}]]
     [[{{node training/SGD/gradients/conv2d_1_1/concat_grad/Slice_1}}]]

Answer 1

问题已解决。如果在定制的Lambda层中使用了tensorflow函数，则需要显式使用set_shape（）函数：

def MyResizeBilinear(x, height, width):
    rows, cols = 1, 2
    original_shape = K.int_shape(x)
    new_shape = tf.constant(np.array([height, width], dtype='int32'))

    x = tf.image.resize_bilinear(x, new_shape, align_corners=True)

    new_height = None if original_shape[rows] is None else height
    new_width = None if original_shape[cols] is None else width

    output_shape = (None, new_height, new_width, None)
    x.set_shape(output_shape)
    return x