How to use these models with Pytorch Lightning

[Boraly]

Hello, I would like to use the models to train them with Pytorch Lightning. So far I have done it by using Automodel.from_pretrained(model_name) in the Lightning class to create the model. Training also works, however I can't get it to load a model from a checkpoint correctly. Also changing the layers will be difficult later on. For example I would like to remove the output layer of a trained model to display the features. What would be the best way to do this?

Thank you, Boraly

[FrancescoSaverioZuppichini]

Could you please share your code? If you don't create your model from lighting then the saved state_dict may be a little different

[Boraly]

So here is my trainingloop:

for model_name in ["resnet152", "densenet201", "vit_base_patch16_224"]:
    model = AutoModel.from_pretrained(model_name)
    model.freeze() 
    if model_name == "densenet201":
      model.head = ResNetHead(in_features=model.encoder.widths[-1], n_classes=4)
    elif model_name == "resnet152":
      model.head = ResNetHead(in_features=model.encoder.widths[-1], n_classes=4)
    else:
      model.head = ViTClassificationHead(n_classes=4)
        # no grad in the encoder
    for param in model.encoder.parameters():
        assert not param.requires_grad
    # grad in the head
    for param in model.head.parameters():
        assert param.requires_grad
    ds = PolypDataModule(data_dir=base, model_name=model_name)
    task = ClassifikationTask(model=model,batch_size=64, lr_rate=0.001, augment="Full", pretrained=True)
    trainer = pl.Trainer(gpus=1, callbacks=[EarlyStopping(monitor="val_loss", patience=8)])
    trainer.fit(task, ds)
    trainer.test(datamodule=ds)

And my pytorch lightning class:

    class ClassifikationTask(pl.LightningModule):
    def __init__(self, model, num_classes=4, batch_size=64, lr_rate=0.001, pretrained=True, augment="Full"):
        super(ClassifikationTask, self).__init__()
        self.model = model
        self.num_classes = num_classes
        self.batch_size = batch_size
        self.model_name = model_name
        self.loss = nn.CrossEntropyLoss()
        self.lr = lr_rate
        self.save_hyperparameters()

        #getting the model
        # if pretrained:
        #   print("PRETRAINED!")

        # else: 
        #   self.model = AutoModel.from_name(self.model_name)

    def training_step(self, batch, batch_idx):
        loss, acc = self._shared_eval_step(batch, batch_idx)
        metrics = {"train_acc": acc, "train_loss": loss}
        self.log_dict(metrics, batch_size=self.batch_size, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        loss, acc = self._shared_eval_step(batch, batch_idx)
        metrics = {"val_acc": acc, "val_loss": loss}
        self.log_dict(metrics, batch_size=self.batch_size, on_epoch=True)
        return loss

    def test_step(self, batch, batch_idx):
        loss, acc = self._shared_eval_step(batch, batch_idx)
        metrics = {"test_acc": acc, "test_loss": loss}
        self.log_dict(metrics, batch_size=self.batch_size )
        return loss

    def _shared_eval_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self.model.forward(x)
        loss = self.loss(y_hat, y)
        acc = FM.accuracy(y_hat, y)
        return loss, acc

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.model.parameters(), lr=self.lr)
        scheduler_dict = {
            "scheduler": ReduceLROnPlateau(
                optimizer,
                factor=0.1,
                patience=5,
                verbose=True
            ),
            "monitor": "val_loss",
            'interval': 'epoch',
        }
        return {"optimizer": optimizer, "lr_scheduler": scheduler_dict}```

Before I also tried to set self.model = AutoModel.from_pretrained(model_name) in the Lightning class, but the results are the same

[FrancescoSaverioZuppichini]

btw to get the features you can

resnet = AutoModel.from_name("resnet18")
resnet.encoder.features
resnet(x)
print(resnet.encoder.features)

The first call to .encoder.features activates the hooks. The code looks fine, could you please point me to the line that is causing the error?

[Boraly]

My problem was a error that I got while I was trying to import a checkpoint. But I resolved my problem. I forgot to save a specific hyperparameter. But I got a new problem. I get an error while trying to import the AutoTransform: File "tSNEVis.py", line 12, in <module> from glasses.models import AutoTransform ImportError: cannot import name 'AutoTransform' from 'glasses.models'

Strangely enough, the import of AutoModel works, which is only one line above. Should I open a new Issue?

[FrancescoSaverioZuppichini]

Hey @Boraly , thank you for the reply. Could you please try to reinstall glasses? I've pushed some changes recently and maybe they'll fix your issue

[Boraly]

Hey, yes this fixed the problem. Thanks!