Hi, I am trying to train a simple LeNet model on Task 0 (that includes class 0 and class 1) save the model path and test its performance with Task 0. Then I load the saved model pth and train it with Task 1 ( that includes class 2 and class 3). When I try to test the model performance with Task 0. It always gives me 0 accuracy.
This is the my entire code
def main():
net = Net().to(device)
train(train_task_0, net, 0)
print(f"Accuracy on test set for TASK 0: {check_accuracy(test_task_0, net) * 100:.2f}")
path = 'model.pt'
net.load_state_dict(torch.load(path))
train(train_task_1, net, 1)
print(f"Accuracy on test set for TASK 1: {check_accuracy(test_task_1, net) * 100:.2f}")
print(f"Accuracy on test set for TASK 0: {check_accuracy(test_task_0, net) * 100:.2f}")
class Net(nn.Module):
def __init__(self, num_classes=10):
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, 40)
self.fc3 = nn.Linear(40, num_classes)
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
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def train(loader, model):
path = 'model.pt'
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
for epoch in range(10): # loop over the dataset multiple
model.train()
running_loss = 0.0
for i, data in enumerate(loader):
optimizer.zero_grad()
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
outputs = model(inputs)
loss = criterion(outputs,labels)
running_loss += loss.item() * inputs.size(0)
loss.backward()
optimizer.step()
torch.save(model.state_dict(), path)
def check_accuracy(loader, model):
num_correct = 0
num_samples = 0
model.eval()
with torch.no_grad():
for x, y in loader:
x = x.to(device=device)
y = y.to(device=device)
print('true labels', y)
scores = model(x)
_, predictions = scores.max(1)
print('predictions',predictions)
num_correct += (predictions == y).sum()
print(num_correct)
num_samples += predictions.size(0)
model.train()
return num_correct/num_samples
if __name__ == '__main__':
main()
This it the output I get:
TRAIN TASK 0, TEST TASK 0
true labels tensor([0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1], device=‘cuda:0’)
predictions tensor([0, 1, 1, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 1, 1, 1], device=‘cuda:0’)
. . .
tensor(1409, device=‘cuda:0’)
Accuracy on test set for TASK 0: 70.45
TRAIN TASK 1, TEST TASK 1
true labels tensor([3, 3, 2, 2, 3, 2, 2, 3, 3, 2, 3, 2, 2, 3, 3, 3], device=‘cuda:0’)
predictions tensor([2, 2, 3, 2, 2, 2, 2, 3, 3, 3, 2, 3, 3, 3, 3, 3], device=‘cuda:0’)
. . .
tensor(1216, device=‘cuda:0’)
Accuracy on test set for TASK 1: 60.80
TRAIN TASK 1, TEST TASK 0
true labels tensor([0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1], device=‘cuda:0’)
predictions tensor([2, 2, 3, 3, 2, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 3], device=‘cuda:0’)
. . .
tensor(0, device=‘cuda:0’)
Accuracy on test set for TASK 0: 0.00
Really appreciated any kind of help, Thanks!