Keras OOM when allocating tensor with shape

Question

I have been scratching my head over this OOM Error for days and I am new to Keras. I have tried sampling down my data, lowering batch size, and removing layers from 3D-Unet but nothing is working for me. I am using LIDC IDRI dataset of CT scans of 1010 Patients. After pre-processing I save my volumes of 64x64x64 shape on disk which I extracted from resampled 256x256x256 whole CT scans (That is because at first I was first trying to train on whole CT scans but after getting OOM I decided to go with 64 cubic size shapes). Each patient has 64 shapes of 64x64x64, and in total that makes 64,640 samples on which I have to train my 3D-Unet.

Here’s my Keras code for the model:

im_width = 64
im_height = 64
im_depth = 64
path_train = 'D:/LIDC-IDRI-Dataset/'

def npz_volume_generator(inputPath, bs, mode="train", aug=None):
    batch_start_index = 0
    patients = os.listdir(inputPath + "images")
    # loop indefinitely
    while True:
        # initialize our batches of scans and masks
        scan_pixels = []
        mask_pixels = []

        # keep looping until we reach our batch size
        for id_ in range(batch_start_index, batch_start_index+bs):
            # attempt to read the next sample from path
            scan_pixel = np.zeros((im_depth, im_width, im_height))
            scan_pixel = np.load(inputPath + 'images/' + patients[id_])['arr_0']
            mask_pixel = np.zeros((im_depth, im_width, im_height))
            mask_pixel = np.load(inputPath + 'masks/' + patients[id_])['arr_0']

            # check to see if we have reached the end of our samples
            if(batch_start_index >= len(patients)):
                # reset the batch start index to the beginning of our samples
                batch_start_index -= len(patients)
                # if we are evaluating we should now break from our
                # loop to ensure we don't continue to fill up the
                # batch from samples from the beginning
                if mode == "eval":
                    break
            # update our corresponding batch lists
            scan_pixels.append(scan_pixel)
            mask_pixels.append(mask_pixel)

        batch_start_index += 1
        if(batch_start_index >= len(patients)):
            batch_start_index -= len(patients)
        # if the data augmentation object is not None, apply it
        if aug is not None:
            (scan_pixels, mask_pixels) = next(aug.flow(np.array(scan_pixels),np.array(mask_pixels), batch_size=bs))

        #Re-shaping and adding a channel dimension (5D Tensor)
        #batch_size, length, breadth, height, channel [None,im_width,im_height,im_depth,1]
        #yield the batch to the calling function
        yield (np.array(expand_dims(scan_pixels, axis=4)), np.array(expand_dims(mask_pixels, axis=4)))


def conv3d_block(input_tensor, n_filters, kernel_size=3, batchnorm=True):
    # first layer
    x = Conv3D(filters=n_filters, kernel_size=(kernel_size, kernel_size, kernel_size), kernel_initializer="he_normal",
               padding="same")(input_tensor)
    if batchnorm:
        x = BatchNormalization()(x)
    x = Activation("relu")(x)
    # second layer
    x = Conv3D(filters=n_filters, kernel_size=(kernel_size, kernel_size, kernel_size), kernel_initializer="he_normal",
               padding="same")(x)
    if batchnorm:
        x = BatchNormalization()(x)
    x = Activation("relu")(x)
    return x

def get_unet(input_img, n_filters=16, dropout=0.5, batchnorm=True):
    # contracting path
    c1 = conv3d_block(input_img, n_filters=n_filters*1, kernel_size=3, batchnorm=batchnorm)
    p1 = MaxPooling3D((2, 2, 2)) (c1)
    p1 = Dropout(dropout*0.5)(p1)

    c2 = conv3d_block(p1, n_filters=n_filters*2, kernel_size=3, batchnorm=batchnorm)
    p2 = MaxPooling3D((2, 2, 2)) (c2)
    p2 = Dropout(dropout)(p2)

    c3 = conv3d_block(p2, n_filters=n_filters*4, kernel_size=3, batchnorm=batchnorm)
    p3 = MaxPooling3D((2, 2, 2)) (c3)
    p3 = Dropout(dropout)(p3)

    c4 = conv3d_block(p3, n_filters=n_filters*16, kernel_size=3, batchnorm=batchnorm)

    # expansive path
    u5 = Conv3DTranspose(n_filters*8, (3, 3, 3), strides=(2, 2, 2), padding='same') (c4)
    u5 = concatenate([u5, c3])
    u5 = Dropout(dropout)(u5)
    c5 = conv3d_block(u5, n_filters=n_filters*8, kernel_size=3, batchnorm=batchnorm)

    u6 = Conv3DTranspose(n_filters*4, (3, 3, 3), strides=(2, 2, 2), padding='same') (c5)
    u6 = concatenate([u6, c2])
    u6 = Dropout(dropout)(u6)
    c6 = conv3d_block(u6, n_filters=n_filters*4, kernel_size=3, batchnorm=batchnorm)

    u7 = Conv3DTranspose(n_filters*2, (3, 3,3), strides=(2, 2, 2), padding='same') (c6)
    u7 = concatenate([u7, c1])
    u7 = Dropout(dropout)(u7)
    c7 = conv3d_block(u7, n_filters=n_filters*2, kernel_size=3, batchnorm=batchnorm)

    outputs = Conv3D(1, (1, 1, 1), activation='sigmoid') (c7)
    model = Model(inputs=[input_img], outputs=[outputs])
    return model

# initialize the number of epochs to train for and batch size
NUM_EPOCHS = 50
BS = 8

# initialize the total number of training and testing image
NUM_TRAIN_IMAGES = len(os.listdir(path_train+ 'images/'))
NUM_TEST_IMAGES = len(os.listdir(path_train+ 'test/'))

# construct the training image generator for data augmentation
aug = ImageDataGenerator(rotation_range=20, zoom_range=0.15,
    width_shift_range=0.2, height_shift_range=0.2, shear_range=0.15,
    horizontal_flip=True, fill_mode="nearest")

# initialize both the training and testing image generators
trainGen = npz_volume_generator(path_train, BS, mode="train", aug=aug)
testGen = npz_volume_generator(path_train, BS, mode="train", aug=None)

# initialize our Keras model and compile it
model = get_unet(Input((im_depth, im_width, im_height, 1)), n_filters=16, dropout=0.05, batchnorm=True)
print(model.summary())
model.compile(optimizer=Adam(), loss="binary_crossentropy", metrics=["accuracy"])

# train the network
print("[INFO] training w/ generator...")
H = model.fit_generator(trainGen, steps_per_epoch=NUM_TRAIN_IMAGES // BS,
                        validation_data=testGen, validation_steps=NUM_TEST_IMAGES // BS,
                        epochs=NUM_EPOCHS)

There are two issues with the output I get. The first warning I get is this:

\Anaconda3\lib\site-packages\keras_preprocessing\image\numpy_array_iterator.py:127: UserWarning: NumpyArrayIterator is set to use the data format convention "channels_last" (channels on axis 3), i.e. expected either 1, 3, or 4 channels on axis 3. However, it was passed an array with shape (8, 64, 64, 64) (64 channels). str(self.x.shape[channels_axis]) + ' channels).')

It states that the shape passed to Keras library was (8, 64, 64, 64) (64 channels), however the input shape I declared in Input() function of Keras is (64, 64, 64, 1) with 1 being the channel on last axis, you don’t declare batch size here which is 8 in my case, yet Keras state that the shape passed on to it has 64 channels, ignoring the last dimension I gave it.

The second error that I get is as following:

ResourceExhaustedError: OOM when allocating tensor with shape[8,32,64,64,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
     [[{{node conv3d_transpose_3/conv3d_transpose}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.
     [[{{node loss/mul}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

Again I have a problem with shape here, Tensor Shape. My shape should be (8,64,64,64,1) but what it reports is (8,32,64,64,64), not only my number of channels are huge here but I also have no idea where that 32 came from. Is there a different interpretation to Tensor Shape? I think there’s something wrong with my input shapes (which is unknowingly being to set to very large) and that is causing the OOM error.

Answer 1

It might be worth switching to half precision floats which will reduce the memory use:

from tensorflow.keras import mixed_precision

policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

And your last layer should be instead:

outputs = Conv3D(1, (1, 1, 1)) (c7)
outputs = layers.Activation('sigmoid', dtype='float32')(outputs)

having your last layer as float16 can lead to instability. You can read more about it here

Answer 2

Add a print statement before your yield statement to check that your actual input actually has a shape of (bs, 64, 64, 64, 1). Defining an Input() of shape (64, 64, 64, 1) just tells your model to expect input of that shape but that does not necessarily mean that your data are generated in that required shape.

The OOM error could be due to having insufficient RAM. The shape of [8,32,64,64,64] is most probably from an intermediate layer and not the input layer. Try lowering the batch size to fix this.