Using Automatic Mixed Precision for existing tochvision models

mixed-precision
1

Hi,
I am trying to use AMP with my resnet50 code from https://github.com/horovod/horovod/blob/master/examples/pytorch/pytorch_imagenet_resnet50.py. I changed the lines 75 through 86 to the following snippet. As I understand from PyTorch documentation, I am supposed to use autocast in the forward pass. But, since I am using a predefined model (torchvision.models.resnet50()), I don’t have access to that part of the code. So, I am getting the error

File “pytorch_imagenet_resnet50_mp.py”, line 78, in train
with autocast:
AttributeError: __enter__

How can I use AMP with these predefined models?

        for i in range(0, len(data), args.batch_size):
            data_batch = data[i:i + args.batch_size]
            target_batch = target[i:i + args.batch_size]
            with autocast:
                output = model(data_batch)
                train_accuracy.update(accuracy(output, target_batch))
                loss = F.cross_entropy(output, target_batch)
                train_loss.update(loss)
                # Average gradients among sub-batches
                loss.div_(math.ceil(float(len(data)) / args.batch_size))
            scaler.scale(loss).backward()
        # Gradient is applied across all ranks
        #optimizer.step()
        scaler.step(optimizer)
2

You would have to call the op via:

with autocast():
    ...
1 Like