# Learning on the fly

Let’s assume that i have trained a cnn model with 5 classes. Now, i want to train the same trained model on 3 new classes. how could i use the trained model to train only the 3 new classes and not on the 8 classes??

Using the tutiorl from Pytorch:
https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html#sphx-glr-beginner-blitz-cifar10-tutorial-py
I think this is a simple question, All you need to do is change the final Linear layer(or fully connected layer) from outputting a matrix that has shape [batch_size, 5] to be [batch_size, 3] (3 new classes)

I suggest that you don’t need to train it again. You can just use transfer learning to only train the last layer which predicts 3 classes only. (you can see the PyTorch tutorial for more detail).

``````before( cnn model with 5 classes):
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 5)     #    5

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x

net = Net()
``````
``````after( cnn model with 3 classes):
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(3, 6, 5)
self.pool = nn.MaxPool2d(2, 2)
self.conv2 = nn.Conv2d(6, 16, 5)
self.fc1 = nn.Linear(16 * 5 * 5, 120)
self.fc2 = nn.Linear(120, 84)
self.fc3 = nn.Linear(84, 3)                 #you said only 3 new classes, you we only predict 3 new class

def forward(self, x):
x = self.pool(F.relu(self.conv1(x)))
x = self.pool(F.relu(self.conv2(x)))
x = x.view(-1, 16 * 5 * 5)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
x = self.fc3(x)
return x

net = Net()
``````

I didn’t mean that. what i mean is i want to train the model on the 3 new classes but the inference predict the whole 8 classes

You can use two different linear layers. You can put a condition like this:

``````if self.train():
self.fc3 = nn.Linear(x,5)
else:
--
``````

That way you can load the weights of this `fc` layer whenever you want.

1 Like