Apply the functions to torchvision model

In this tutorial, we are going to apply rational functions to a (pretrained) PyTorch network directly taken from the torchvision library.

First, we import the torchvision library:

import torchvision

Then let’s instantiate a torchvision model on the first cuda-enabled GPU:

device = 'cuda:0'
vgg16 = torchvision.models.vgg16(pretrained=False).cuda(device)

Now we can take the VGG-16 model and convert all its ReLU layers to Rationals with the help of a utility function:

from rational.utils.convert_network import convert_pytorch_model_to_rational
model = convert_pytorch_model_to_rational(vgg16, rational_cuda=True)

The function will scan the network for certain activation functions and will replace them with rationals that are initialized to approximate the same function.