Fine tune a new dataset with pre-trained Imagenet weights SimCLR

[jmarrietar]

Hi Everyone!. I am new to Pytorch Lighting and would like to use the pre-trained Simclr model weights with Imagenet and Fine-tune it using another dataset.

I try first to fine-tune it using CIFAR-10, but I am getting the following error, not sure what I am doing wrong or what should I change in my code.

TypeError: conv2d(): argument 'input' (position 1) must be Tensor, not list

Any Help is very Appreciate it

Colab: https://colab.research.google.com/drive/1V-NLNUH8ltaPwnJQdDDPRdQofm7XKylv?usp=sharing

My code:

import pytorch_lightning as pl
from pl_bolts.models.self_supervised import SimCLR
from pl_bolts.datamodules import CIFAR10DataModule
from pl_bolts.models.self_supervised.simclr.transforms import (
    SimCLREvalDataTransform, SimCLRTrainDataTransform)
from pl_bolts.models.self_supervised.ssl_finetuner import SSLFineTuner

#  Data from Cifar10
dm = CIFAR10DataModule(num_workers=0)
dm.train_transforms = SimCLRTrainDataTransform(32)
dm.val_transforms = SimCLREvalDataTransform(32)

# Backbone weigths from ImageNet
from pl_bolts.models.self_supervised import SimCLR

weight_path = 'https://pl-bolts-weights.s3.us-east-2.amazonaws.com/simclr/bolts_simclr_imagenet/simclr_imagenet.ckpt'
backbone = SimCLR(gpus=1, num_samples=dm.num_samples, batch_size=dm.batch_size, dataset='imagenet').load_from_checkpoint(weight_path, strict=False)

tuner = SSLFineTuner(
    backbone,
    in_features=backbone.feat_dim,
    num_classes=10,
    epochs=100,
    hidden_dim=1024,
    dropout=0.,
    learning_rate=0.3,
    weight_decay=1e-6,
    nesterov=False,
    scheduler_type='cosine',
    gamma=0.1,
    final_lr=0.
)

trainer = pl.Trainer(
    gpus=1,
    num_nodes=1,
    precision=16,
    max_epochs=100
)

trainer.fit(tuner, dm)