Loading and Training a Neural Network with Custom dataset via Transfer Learning in Pytorch
7 min readJul 16, 2018
In the previous post (here), we loaded and transformed custom images from a directory of training and validation datasets into appropriately processed Tensors; now we are ready to load, modify, train and test an existing model with our readymade data, in four steps:
- Loading a Neural Network model
- Building the classifier and training the network
- Testing the modified network
- Saving the checkpoint
Loading the Neural Network model
There are a variety of existing Neural Networks(NN), trained on vast amounts of datasets such as Imagenet, Kaggle and the UCI repository just to state a few. The graph below describes such public NN models, on a scale of the accuracy achieved upon conception, with respect to the dataset size used for training.
To load the NN model of a preferred type, import the ‘models’ package from ‘torchvision’ and call your desired model with the required parameters:
#import models from torchvision
from torchvision import models#build the pretrained model (vgg16 in this case)
model = models.vgg16(pretrained = True)
