Блог

Я пытаюсь заново обучить верхние уровни SSD-сети, адаптировав учебник по обнаружению объектов с несколькими кадрами.

Обучение последовательно завершается неудачно после многих партий, но до завершения с ошибкой

 Allocator (GPU_0_bfc) ran out of memory trying to allocate 1.56GiB (rounded to 1677721600)requested by op ResNet50V1_FPN/model/conv2_block2_3_bn/FusedBatchNormV3
If the cause is memory fragmentation maybe the environment variable 'TF_GPU_ALLOCATOR=cuda_malloc_async' will improve the situation. 
Current allocation summary follows.
Current allocation summary follows.
2021-07-19 11:51:56.505408: I tensorflow/core/common_runtime/bfc_allocator.cc:991] BFCAllocator dump for GPU_0_bfc
2021-07-19 11:51:56.505426: I tensorflow/core/common_runtime/bfc_allocator.cc:998] Bin (256):   Total Chunks: 27452, Chunks in use: 27452. 6.70MiB allocated for chunks. 6.70MiB in use in bin. 1.47MiB client-requested in use in bin.
...
 

Я тренируюсь на V100 с 16 ГБ памяти.

Мой код создает файл tf.data.Набор данных пар (байты изображения, идентификатор изображения). Изображения загружаются с помощью tf.io методы

 def load_image_tf(path):
    """Load an image from file to TensorFlow tensor.
    
    Args:
        path: path of file
        
    Returns:
        Tensor: image data tensor
    """
    
    img_bytes = tf.io.read_file(path)
    
    img_tensor = tf.io.decode_jpeg(img_bytes, channels=3)
    
    img_tensor = tf.image.resize(img_tensor, [640,640])
    
    return img_tensor
 

Ограничительные рамки и классы Groundtruth считываются из файла .csv в (N,4) и одномерные тензоры, а затем в словарь, индексируемый идентификатором изображения перед обучением.

Мой цикл обучения считывает набор данных пакетами и использует идентификаторы изображений из пакета для индексирования словарей основной истины.

 for batch in train_ds:
    
    batch_keys = batch[1]
    gt_boxes_list = [train_bboxes[key.numpy().decode()]['bboxes'] for key in batch_keys]
    gt_classes_list = [train_bboxes[key.numpy().decode()]['classes'] for key in batch_keys]
    
    image_tensors = batch[0]
    
    # Training step (forward pass   backward pass)
    total_loss = train_step_fn(image_tensors, gt_boxes_list, gt_classes_list)
 

Списки ограничивающих прямоугольников и классов отправляются в функцию train.

     @tf.function
    def train_step_fn(image_tensors,
                     groundtruth_boxes_list,
                     groundtruth_classes_list):
        """A single training iteration.
        
        Args:
            image_tensors: A list of [1, height, width, 3] Tensor of type tf.float32.
                Note that the height and width can vary across images, as they are
                reshaped within this function to be 640x640.
            groundtruth_boxes_list: A list of Tensors of shape [N_i, 4] with type
                tf.float32 representing groundtruth boxes for each image in the batch.
            groundtruth_classes_list: A list of Tensors of shape [N_i, num_classes]
                with type tf.float32 representing groundtruth boxes for each image in
                the batch.

        Returns:
            A scalar tensor representing the total loss for the input batch.
        """
        
        #shapes = tf.constant(batch_size * [[800, 600, 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)
            
            preprocessed_images, shapes = model.preprocess(image_tensors)
            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
 

Я подозреваю, что я делаю что-то вроде создания нового тензора на каждой итерации, который не очищается, но я не вижу, где именно это происходит или какая у меня альтернатива. Я не могу включить основную истину в объект набора данных, так как на каждом изображении имеется переменное количество ограничивающих прямоугольников, в результате чего каждый тензор основной истины имеет различную форму.