tensorflow improper batch setup Check failed: cudnnSetTensorNdDescriptor when using upsampling - tensorflow

Question

I think I am setting up my batches wrong. If I run with the generated dataset it runs fine but with my own data I get an error.
If I take out the encoder (max pooling) and decoder (UpSampling2D) I don't get an error.
input size (304, 228, 1)
Generated: RUNS
import tensorflow as tf
from tensorflow.keras import layers
from natsort import natsorted
from tensorflow.keras.models import Model
BATCH_SIZE = 4
EPOCHS = 20
LEARNING_RATE = 1e-4
RESET_TRAINING = True
INPUT_CHANNELS = 1
OUTPUT_CHANNELS = 1
LOSS_TYPE = tf.keras.losses.SparseCategoricalCrossentropy()
img_size = (304, 228)
# configure cuda
physical_devices = tf.config.experimental.list_physical_devices('GPU')
assert len(physical_devices) > 0, "Not enough GPU hardware devices available"
config = tf.config.experimental.set_memory_growth(physical_devices[0], True)
class UnetModel(Model):
def __init__(self, *args, **kwargs):
super().__init__(UnetModel, *args, **kwargs)
# -- Encoder -- #
# Block encoder 1
input_shape = (img_size[0], img_size[1], 1)
# If you want to know more about why we are using `he_normal`:
# https://stats.stackexchange.com/questions/319323/whats-the-difference-between-variance-scaling-initializer-and-xavier-initialize/319849#319849
# Or the excellent fastai course:
# https://github.com/fastai/course-v3/blob/master/nbs/dl2/02b_initializing.ipynb
initializer = 'he_normal'
inputs = layers.Input(shape=input_shape)
print("input shape ", input_shape)
conv_enc_1 = layers.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer=initializer)(inputs)
conv_enc_1 = layers.Conv2D(64, 3, activation = 'relu', padding='same', kernel_initializer=initializer)(conv_enc_1)
# Block encoder 2
max_pool_enc_2 = layers.MaxPooling2D(pool_size=(2, 2))(conv_enc_1)
conv_enc_2 = layers.Conv2D(128, 5, activation = 'relu', padding = 'same', kernel_initializer = initializer)(max_pool_enc_2)
conv_enc_2 = layers.Conv2D(128, 5, activation = 'relu', padding = 'same', kernel_initializer = initializer)(conv_enc_2)
# Block decoder 1
up_dec_4 = layers.Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = initializer)(layers.UpSampling2D(size = (2,2))(conv_enc_2))
merge_dec_4 = layers.concatenate([conv_enc_1, up_dec_4], axis = 3)
conv_dec_4 = layers.Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = initializer)(merge_dec_4)
conv_dec_4 = layers.Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = initializer)(conv_dec_4)
conv_dec_4 = layers.Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = initializer)(conv_dec_4)
# -- Dencoder -- #
output = layers.Conv2D(1, 1, activation = 'softmax')(conv_dec_4)
self.model = tf.keras.Model(inputs = inputs, outputs = output)
def call(self, x):
return self.model(x)
model = UnetModel()
model.compile(optimizer=tf.keras.optimizers.Adam(LEARNING_RATE), loss = LOSS_TYPE, metrics= [tf.keras.metrics.get('accuracy')])
dataset_debug = tf.data.Dataset.from_tensor_slices((tf.random.normal(shape = (BATCH_SIZE, img_size[0], img_size[1], 1)), tf.random.normal(shape = (BATCH_SIZE, img_size[0], img_size[1], 1)))).batch(BATCH_SIZE)
history = model.fit(dataset_debug, epochs=EPOCHS, shuffle=True)
Does NOT run
Here I am splitting the filenames into training and validation sets using train_test_split and reading in images in the parse_img_input function
# takes image filenames of uint8 and normalizes to 0-1 range
def parse_img_input(img_file, img_file_out):
print("img file ", img_file)
def _parse_input(img_file, img_file_out):
# get img image
d_filepath = img_file.numpy().decode()
d_image_decoded = tf.image.decode_jpeg(
tf.io.read_file(d_filepath), channels=1)
d_image = tf.cast(d_image_decoded, tf.float32) / 255.0
# get img image
d_filepath_out = img_file_out.numpy().decode()
d_image_decoded_out = tf.image.decode_jpeg(
tf.io.read_file(d_filepath_out), channels=1)
d_image_out = tf.cast(d_image_decoded_out, tf.float32) / 255.0
# add channel dimension
d_image = tf.expand_dims(d_image, -1)
d_image_out = tf.expand_dims(d_image_out, -1)
return d_image, d_image_out
return tf.py_function(_parse_input,
inp=[img_file, img_file_out],
Tout=[tf.float32, tf.float32])
# depth_files_in, depth_files_out are lists of filenames
# split input data into train, test sets
X_train_file, X_test_file, y_train_file, y_test_file = train_test_split(depth_files_in, depth_files_out,
test_size=0.2,
random_state=0)
dataset_train = tf.data.Dataset.from_tensor_slices((X_train_file, y_train_file))
dataset_train = dataset_train.map(parse_img_input)
dataset_test = tf.data.Dataset.from_tensor_slices((X_test_file, y_test_file))
dataset_test = dataset_test.map(parse_img_input)
history = model.fit(dataset_train, epochs=EPOCHS, shuffle=True, batch_size = BATCH_SIZE, validation_data= dataset_test)
F tensorflow/stream_executor/cuda/cuda_dnn.cc:535] Check failed: cudnnSetTensorNdDescriptor(handle_.get(), elem_type, nd, dims.data(), strides.data()) == CUDNN_STATUS_SUCCESS (3 vs. 0)batch_descriptor: {count: 228 feature_map_count: 64 spatial: 152 0 value_min: 0.000000 value_max: 0.000000 layout: BatchDepthYX}