discussion: Can we use our own Pytorch models in avalanche stragtegy?

[akashgupta1909]

I have been trying to run my own PyTorch model, which is a LightningModule Derivative, using benchmarking experience which i got from avalanche. But it gives the following error:

Training on experience 0
-- >> Start of training phase << --
Traceback (most recent call last):
  File "/Users/akashgupta/Documents/Projects/Python/binary_deepfake_detection/model.py", line 387, in <module>
    strategy.train(experience)  # Train on the current experience
    ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/avalanche/training/templates/base_sgd.py", line 211, in train
    super().train(experiences, eval_streams, **kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/avalanche/training/templates/base.py", line 162, in train
    self._before_training_exp(**kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/avalanche/training/templates/base_sgd.py", line 291, in _before_training_exp
    self.make_train_dataloader(**kwargs)
  File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/avalanche/training/templates/base_sgd.py", line 456, in make_train_dataloader
    self.dataloader = TaskBalancedDataLoader(
                      ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Library/Frameworks/Python.framework/Versions/3.12/lib/python3.12/site-packages/avalanche/benchmarks/utils/data_loader.py", line 404, in __init__
    task_labels_field = getattr(data, "targets_task_labels")
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AttributeError: 'ClassificationDataset' object has no attribute 'targets_task_labels'
(base) MacBook-Pro:binary_deepfake_detection akashgupta$ ;2B

The code which I've been using:

    train_dataset_wrapped = as_classification_dataset(train_dataset)
    test_dataset_wrapped = as_classification_dataset(val_dataset)
    train_experiences = [train_dataset_wrapped]  # List of training experiences
    test_experiences = [test_dataset_wrapped]    # List of testing experiences

    # Create a benchmark scenario from the experiences
    benchmark = benchmark_from_datasets(
        train=train_experiences,
        test=test_experiences
    )

    # Initialize model, optimizer, and criterion
    model = BNext4DFR(num_classes=2)  # Your model
    optimizer = Adam(model.parameters())
    criterion = nn.CrossEntropyLoss()

    # Setup the strategy
    strategy = Naive(
        model=model,
        optimizer=optimizer,
        criterion=criterion
    )

    # Training loop
    for experience in benchmark.train_stream:
        print(f"Training on experience {experience.current_experience}")
        strategy.train(experience)  # Train on the current experience

    # Evaluation loop
    for experience in benchmark.test_stream:
        print(f"Testing on experience {experience.current_experience}")
        strategy.eval(experience)

Can someone suggest how to make it work?