2

I built my network the following way: 

# Build U-Net model
inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
s = Lambda(lambda x: x / 255) (inputs)
width = 64
c1 = Conv2D(width, (3, 3), activation='relu', padding='same') (s)
c1 = Conv2D(width, (3, 3), activation='relu', padding='same') (c1)
p1 = MaxPooling2D((2, 2)) (c1)

c2 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (p1)
c2 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (c2)
p2 = MaxPooling2D((2, 2)) (c2)

c3 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (p2)
c3 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (c3)
p3 = MaxPooling2D((2, 2)) (c3)

c4 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (p3)
c4 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (c4)
p4 = MaxPooling2D(pool_size=(2, 2)) (c4)

c5 = Conv2D(width*16, (3, 3), activation='relu', padding='same') (p4)
c5 = Conv2D(width*16, (3, 3), activation='relu', padding='same') (c5)

u6 = Conv2DTranspose(width*8, (2, 2), strides=(2, 2), padding='same') (c5)
u6 = concatenate([u6, c4])
c6 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (u6)
c6 = Conv2D(width*8, (3, 3), activation='relu', padding='same') (c6)

u7 = Conv2DTranspose(width*4, (2, 2), strides=(2, 2), padding='same') (c6)
u7 = concatenate([u7, c3])
c7 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (u7)
c7 = Conv2D(width*4, (3, 3), activation='relu', padding='same') (c7)

u8 = Conv2DTranspose(width*2, (2, 2), strides=(2, 2), padding='same') (c7)
u8 = concatenate([u8, c2])
c8 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (u8)
c8 = Conv2D(width*2, (3, 3), activation='relu', padding='same') (c8)

u9 = Conv2DTranspose(width, (2, 2), strides=(2, 2), padding='same') (c8)
u9 = concatenate([u9, c1], axis=3)
c9 = Conv2D(width, (3, 3), activation='relu', padding='same') (u9)
c9 = Conv2D(width, (3, 3), activation='relu', padding='same') (c9)

outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)
with tf.device('/cpu:0'):
    model = Model(inputs=[inputs], outputs=[outputs])

sgd = optimizers.SGD(lr=0.03, decay=1e-6, momentum=0.9, nesterov=True)
parallel_model = multi_gpu_model(model, gpus=8)
parallel_model.compile(optimizer=sgd, loss='binary_crossentropy', metrics=[mean_iou])
model.summary()

Notice that I am instantiating the base model on the CPU as suggested by keras documentation. Then, I run the network using the following lines: 

# Fit model
earlystopper = EarlyStopping(patience=20, verbose=1)
checkpointer = ModelCheckpoint('test.h5', verbose=1, save_best_only=True)
results = parallel_model.fit(X_train, Y_train, validation_split=0.05, batch_size = 256, verbose=1, epochs=100, 
                    callbacks=[earlystopper, checkpointer])

However, even though, I am using the multiple_gpu_model, my code still results in the following error:

OOM when allocating tensor with shape[32,128,256,256] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc

Which indicates that the network is trying to run the batch size of 256 on just a single GPU instead of 8. Am I not implementing this properly? Do I need to use Xception as in the example?

Jonathan
  • 1,876
  • 2
  • 20
  • 56

1 Answers1

2

The first dim of the tensor is the batch_size, so everthing is fine in your case. You have specified your batch_size as 256 and you use 8 gpus. So your resulting batch_size is 32 as stated in the error. Also the error suggest that your model still is too big with a batch_size of 32 for your gpus to handle.

dennis-w
  • 2,166
  • 1
  • 13
  • 23
  • 1
    I am able to run a batch size of 32 on a single gpu. So does the memory usage change when gpu usage changes? – Jonathan Apr 26 '18 at 07:25
  • Seems so, this is interesting. My guess would be that tensorflow need additional GPU memory for sharing the weights between GPUs. How much free gpu space do you have when training with a batch_size of 32 on a single GPU? – dennis-w Apr 26 '18 at 07:31
  • Bringing it down to 128 did fix the problem, but now there's a new problem: https://stackoverflow.com/q/50243086/8100895 – Jonathan May 08 '18 at 22:26