I’m currently training a retinanet model for fashion apparel with pytorch. However I’m getting very slow training times (0.5s per step on 1333x1333 inputs).
I’m training this on a queue server which makes debugging very hard (also I have a job scheduling limit).
Can anyone help me with a code review?
The main bits I need reviews with is at collate_fn and with the hist object.
import torch

from torch_collections.models.retinanet import RetinaNet
from torch_collections.losses import DetectionFocalLoss, DetectionSmoothL1Loss

def main():
    # Load dataset
    #### INSERT YOUR DATASET HERE ####

    # dataset is a torch.utils.data.Dataset object
    # dataset should output samples of dict type with keys 'image' and 'annotations'
    # Here is what a sample would look like
    # sample = {
    #     'image'       : A numpy.ndarray image of dtype uint8 in RGB, HWC format and
    #     'annotations' : A numpy.ndarray of shape (number_of_annotations, 5)
    #                     Each annotation is of the format (x1, y1, x2, y2, class_id)
    # }
    dataset = YOUR_DETECTION_DATASET

    ##################################

    # Load model with initial weights
    #### INSERT NUMBER OF CLASSES HERE ####
    retinanet = RetinaNet(num_classes=NUMBER_OF_CLASSES)
    #######################################

    # Create dataset iterator
    dataset_loader = torch.utils.data.DataLoader(
        dataset,
        batch_size=1,
        shuffle=False,
        collate_fn=retinanet.collate_fn,
        num_workers=2,
    )
    # The retinanet.collate_fn uses GPU to perform a rather costly transformation
    # Some tensors here are shared, my reasoning for using 2 workers is to read images
    # from files faster (some images are high-res)
    # This could possibly be a major road block as there could be needless IO from GPU to CPU
    # and back.

    # Initialize loss functions and optimizer
    focal_loss_fn = DetectionFocalLoss()
    huber_loss_fn = DetectionSmoothL1Loss()
    optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, retinanet.parameters()), lr=0.00001)

    for epoch in range(50):
        dataset_iterator = dataset_loader.__iter__()
        hist = {
            'loss': [],
            'regression_loss': [],
            'classification_loss': []
        }

        for step_nb in range(10000):
            # Get sample and zero optimizer
            sample = dataset_iterator.next()
            optimizer.zero_grad()

            # forward
            outputs = retinanet(sample['image'])
            regression_loss = huber_loss_fn(outputs[0], sample['regression'])
            classification_loss = focal_loss_fn(outputs[1], sample['classification'])

            # If all anchors are negative then we just skip this epoch
            if classification_loss is None:
                continue

            # If there are no positive anchors we skip regression loss
            if regression_loss is None:
                regression_loss = 0

            loss = regression_loss + classification_loss

            # backward + optimize
            loss.backward()
            optimizer.step()

            # *** My main question is if storing step data like this counts as making
            # cuda copies (which I heard is an expensive process)
            hist['loss'].append(loss)
            hist['regression_loss'].append(regression_loss)
            hist['classification_loss'].append(classification_loss)

            # Clean batch
            # *** I am not sure if this is a required step but somehow I stopped
            # getting cuda runtime error (30) after this
            del loss
            del regression_loss
            del classification_loss
            del sample['image']
            del sample['regression']
            del sample['classification']
            del sample

        torch.save(retinanet, 'snapshot/epoch_{}.pth'.format(epoch))
        print('Epoch {} - loss: {} - regression: {} - classification: {}'.format(
            epoch,
            torch.sum(torch.stack(hist['loss'])).item(),
            torch.sum(torch.stack(hist['regression_loss'])).item(),
            torch.sum(torch.stack(hist['classification_loss'])).item()
        ))

    print('Finished Training')


if __name__ == '__main__':
    # This step became required for python 3.X to share tensors for multiprocessing
    torch.multiprocessing.set_start_method("spawn", force=True)
    main()