Graph retracing consuming all memory with no workaround

The below is my chunk of code--> very similar to the rubber ducky example I have commented on below.

tf.keras.backend.set_learning_phase(True)

# These parameters can be tuned; since our training set has 5 images
# it doesn't make sense to have a much larger batch size, though we could
# fit more examples in memory if we wanted to.
batch_size = 16
learning_rate = 0.01
num_batches = 108

# Select variables in top layers to fine-tune.
trainable_variables = detection_model.trainable_variables
to_fine_tune = []
prefixes_to_train = [
  'WeightSharedConvolutionalBoxPredictor/WeightSharedConvolutionalBoxHead',
  'WeightSharedConvolutionalBoxPredictor/WeightSharedConvolutionalClassHead']
for var in trainable_variables:
  if any([var.name.startswith(prefix) for prefix in prefixes_to_train]):
    to_fine_tune.append(var)

# Set up forward + backward pass for a single train step.
def get_model_train_step_function(model, optimizer, vars_to_fine_tune):
  """Get a tf.function for training step."""

  # Use tf.function for a bit of speed.
  # Comment out the tf.function decorator if you want the inside of the
  # function to run eagerly.
  @tf.function(experimental_relax_shapes=True)#(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32), # These input sigtensors cause a graph error
#                      tf.TensorSpec(shape=(None,4), dtype=tf.float32),
#                      tf.TensorSpec(shape=None, dtype=tf.float32)]) #(experimental_relax_shapes=True)
  def train_step_fn(image_tensors,
                    groundtruth_boxes_list,
                    groundtruth_classes_list):
   
    shapes = tf.constant(batch_size * [[1104, 1280, 3]], dtype=tf.int32)
    model.provide_groundtruth(
        groundtruth_boxes_list=groundtruth_boxes_list,
        groundtruth_classes_list=groundtruth_classes_list)
    with tf.GradientTape() as tape:
      preprocessed_images = tf.concat(
          [detection_model.preprocess(image_tensor)[0]
           for image_tensor in image_tensors], axis=0)
      prediction_dict = model.predict(preprocessed_images, shapes)
      losses_dict = model.loss(prediction_dict, shapes)
      total_loss = losses_dict['Loss/localization_loss'] + losses_dict['Loss/classification_loss']
      gradients = tape.gradient(total_loss, vars_to_fine_tune)
      optimizer.apply_gradients(zip(gradients, vars_to_fine_tune))
    return total_loss

  return train_step_fn
    
optimizer = tf.keras.optimizers.SGD(learning_rate=learning_rate, momentum=0.9)
train_step_fn = get_model_train_step_function(
    detection_model, optimizer, to_fine_tune)

print('Start fine-tuning!', flush=True)
for idx in range(num_batches):
  # Grab keys for a random subset of examples
  all_keys = list(range(len(train_images_np)))
  random.shuffle(all_keys)
  example_keys = all_keys[:batch_size]

  # Note that we do not do data augmentation in this demo.  If you want a
  # a fun exercise, we recommend experimenting with random horizontal flipping
  # and random cropping :)

  gt_boxes_list = [gt_box_tensors[key] for key in example_keys]
  gt_classes_list = [gt_classes_one_hot_tensors[key] for key in example_keys]
  image_tensors = [train_image_tensors[key] for key in example_keys]

    
  # Training step (forward pass + backwards pass)
  total_loss = train_step_fn(image_tensors, gt_boxes_list, gt_classes_list)

  
  print('batch ' + str(idx) + ' of ' + str(num_batches) + ', loss=' +  str(total_loss), flush=True)

print('Done fine-tuning!')

2. Describe the bug

My code is above - this is nearly identical to the eager_few_shot_od_training_tf2_colab.ipynb

The only difference between my project and this one is I have multiple bounding boxes per image, 1 to 6 per image.

I am running on a Tesla V100 32GB GPU.

Originally when I ran the model with no modification to the @tf.function, I was receiving OOM error due to what I believe is the graph retracing done by this function. In the eager_few_shot_od_training_tf2_colab.ipynb example - since there is only one bounding box per image, this train_step_fn always had the same shape arguments since each gt_box_tensor was of shape (1,4). Well now since I have (1,4),(2,4),(3,4),(4,4),(5,4),(6,4), the @tf.function is retracing too many graphs and causing OOM.

The work arounds I attempted are

@tf.function(experimental_relax_shapes=True) 

# which caused -->

NotImplementedError: Cannot convert a symbolic Tensor (Loss/strided_slice_3:0) to a numpy array. This error may indicate that you're trying to pass a Tensor to a NumPy call, which is not supported

and

@tf.function(input_signature=[tf.TensorSpec(shape=None, dtype=tf.float32), 
                     tf.TensorSpec(shape=None, dtype=tf.float32),
                      tf.TensorSpec(shape=None, dtype=tf.float32)])

which led to -

OperatorNotAllowedInGraphError: iterating over `tf.Tensor` is not allowed: AutoGraph did convert this function. This might indicate you are trying to use an unsupported feature.

I am unaware of the next way to proceed with these errors.

I am using:

Python 3.7
Tensorflow 2.5.0
Numpy 1.20.3
cudnn 8.1.0.77
cudatoolkit 11.2.2