Class Indices Overwritten During Continual Learning with Avalanche 0.4.0a

[Funkvay]

I'm encountering an issue with class indices being overwritten when performing continual learning using Avalanche 0.4.0a.

Environment:

Avalanche Version: 0.4.0a
OS: Ubuntu 20.04
Python Environment: default (not using Conda)

I'm employing Avalanche for continual learning with the aim that when the model trains on new data, it retains knowledge of the old data and successfully integrates the new information. Despite the dataset size, this indexing issue persists. I've been using a small dataset of approximately 10-15 photos for simplicity in troubleshooting. Specifically, after adding 3 new classes to a model initially trained on 5 classes (indexed 0-4), I expected predictions to be indexed as 0,1,2,3,4,5,6,7. However, I received predictions as 0,1,2,3,4,0,1,2.

Here is the code that I have:

import argparse
import os
from torchvision.models import resnet101
import torch
from torch.nn import CrossEntropyLoss
from torchvision import transforms, datasets
from avalanche.benchmarks import nc_benchmark
from avalanche.training.supervised import Naive
from avalanche.training.plugins import ReplayPlugin
from avalanche.evaluation.metrics import (
    forgetting_metrics,
    accuracy_metrics,
    loss_metrics,
)
from avalanche.logging import InteractiveLogger
from avalanche.training.plugins import EvaluationPlugin
import torch.nn.init as init


def main(args):
    # Load the saved model's weights
    saved_weights = torch.load(args.model_path)
    num_old_classes = saved_weights['fc.weight'].shape[0]

    # Initialize the ResNet-101 model
    model = resnet101()

    # Adjust the final layer to match the number of classes in the saved model
    num_features = model.fc.in_features
    model.fc = torch.nn.Linear(num_features, num_old_classes)
    model.load_state_dict(saved_weights)

    # Determine the number of classes in the new dataset
    num_new_classes = len(os.listdir(os.path.join(args.data_path, 'train')))
    total_classes = num_old_classes + num_new_classes
    
    # Expand the final layer of the model to accommodate the new classes
    weights = model.fc.weight.data
    biases = model.fc.bias.data
    model.fc = torch.nn.Linear(in_features=num_features, out_features=total_classes)
    with torch.no_grad():
        model.fc.weight[:num_old_classes] = weights
        model.fc.bias[:num_old_classes] = biases

    # Initialize weights for the new classes
    init.kaiming_uniform_(model.fc.weight[num_old_classes:], mode='fan_in', nonlinearity='relu')
    init.zeros_(model.fc.bias[num_old_classes:])

    model.train()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Define data transformations
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])

    # Load datasets
    train_dataset = datasets.ImageFolder(root=os.path.join(args.data_path, 'train'), transform=transform)
    test_dataset = datasets.ImageFolder(root=os.path.join(args.data_path, 'val'), transform=transform)

    # Adjust labels for the new dataset to be continuous with the old dataset
    train_dataset.targets = [label + num_old_classes for label in train_dataset.targets]
    test_dataset.targets = [label + num_old_classes for label in test_dataset.targets]

    # Create a continual learning benchmark
    benchmark = nc_benchmark(train_dataset, test_dataset, n_experiences=num_new_classes, task_labels=False, seed=1234)

    # Define logging and evaluation plugins
    interactive_logger = InteractiveLogger()
    eval_plugin = EvaluationPlugin(
        accuracy_metrics(minibatch=True, epoch=True, experience=True, stream=True),
        loss_metrics(minibatch=True, epoch=True, experience=True, stream=True),
        forgetting_metrics(experience=True),
        loggers=[interactive_logger],
    )

    # Define the continual learning strategy
    cl_strategy = Naive(
        model,
        torch.optim.SGD(model.parameters(), lr=0.01),
        CrossEntropyLoss(),
        train_mb_size=8,
        train_epochs=10,
        eval_mb_size=8,
        device=device,
        plugins=[ReplayPlugin(mem_size=1005)],
        evaluator=eval_plugin,
    )

    # Continual learning training loop
    for experience in benchmark.train_stream:
        cl_strategy.train(experience)

    # Save the updated model
    torch.save(model.state_dict(), args.model_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model_path", type=str, required=True, help="Path to the pretrained model.")
    parser.add_argument("--data_path", type=str, required=True, help="Path to the new dataset directory.")
    args = parser.parse_args()
    main(args)

I'm trying to understand how to preserve accurate class indexing throughout retraining and evaluation, ideally without changing the core architecture and using Avalanche.

Please suggest how to solve this and make indexation right.

[AlbinSou]

I have tried a similar example, using CIFAR10 instead of a custom dataset, and I did not have any problem. I don't think the problem comes from the nc_benchmark function. In this case, the targets of experience one are [10, 11]. Could you try to print the classes_in_this_experience attribute of the first experience and give the output ?

import torchvision
import torch
import numpy as np

import avalanche

train_ds = torchvision.datasets.CIFAR10(root="/home/albin/Datasets/", train=True)
test_ds = torchvision.datasets.CIFAR10(root="/home/albin/Datasets/", train=True)

transforms = torchvision.transforms.ToTensor()

train_ds.targets = np.array([label + 10 for label in train_ds.targets])
test_ds.targets = np.array([label + 10 for label in test_ds.targets])


benchmark = avalanche.benchmarks.nc_benchmark(train_ds, test_ds, n_experiences=5, task_labels=False, fixed_class_order=np.arange(10, 20), train_transform=transforms, eval_transform=transforms)

experience = benchmark.train_stream[0]
print(experience.classes_in_this_experience)

[Funkvay]

Thank you for your suggestions.
I've tried executing the below-modified version of the code for loading a pretrained model and also modifying the labels of CIFAR10 dataset and creating a benchmark instance.Here I also got [10, 11] :

import torch
import torchvision
import numpy as np
import avalanche

# Root directory for my dataset.
root = './data'

transforms = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

model = torchvision.models.resnet101(weights=torchvision.models.resnet.ResNet101_Weights.IMAGENET1K_V1)

train_ds = torchvision.datasets.CIFAR10(root=root, train=True,  transform=transforms)
test_ds = torchvision.datasets.CIFAR10(root=root, train=False,  transform=transforms)

train_ds.targets = [label + 10 for label in train_ds.targets]
test_ds.targets = [label + 10 for label in test_ds.targets]

benchmark = avalanche.benchmarks.nc_benchmark(train_ds, test_ds, n_experiences=5, task_labels=False, fixed_class_order=np.arange(10, 20), train_transform=transforms, eval_transform=transforms)

experience = benchmark.train_stream[0]
print(experience.classes_in_this_experience)

But when I tried to test the model with this script

import torch
import torchvision.transforms as transforms
from PIL import Image
import torch.nn.functional as F
import torchvision

transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

model = torchvision.models.resnet101()
model.load_state_dict(torch.load("saved_model.pth"))
model.eval()  # Set the model to evaluation mode

image_path = "image.jpg"
image = Image.open(image_path)
input_image = transform(image).unsqueeze(0)


with torch.no_grad():
    output = model(input_image)

# Get the predicted label
probabilities = F.softmax(output, dim=1)
confidence, predicted_class = torch.max(probabilities, dim=1)
print(f"The model is {confidence.item()*100:.2f}% confident that the given image belongs to class {predicted_class.item()}.")

and here I got The model is 95.34% confident that the given image belongs to class 404. when used airplane for testing, which I think is not right because the CIFAR10 should not have 404 class and also the resnet101 has only animals in it (there is not airplane as far as I know). So I'm not sure if this works right or no. Maybe my testing script is not right?

I also have followed your advice and added (in my original code where I actually had the main problem) the respective code to print the classes_in_this_experience of the first experience in the training stream of my benchmark, as shown below:

    for experience in benchmark.train_stream:
        cl_strategy.train(experience)
        results.append(cl_strategy.eval(benchmark.test_stream))
        
     first_experience = benchmark.train_stream[0]
     print(first_experience.classes_in_this_experience)

This prints [1002]. To give a bit more context, I am using ResNet-101, and there are 1000 classes, indexed from 0 to 999. Now, I’ve added 6 new classes.

So right now I'm a little bit confused and don't know which direction to move.