Feature request: Simpler loading of custom models

[dekuenstle]

Using your toolbox with the built-in models is straightforward, but we would like to compare some custom pytorch models.
It would be great to have a routine to add these models (i.e. subclasses of nn.Module) to the toolbox registry from your own script. If this is already possible, it would be great if you could share an example.

Currently, we add the model inside the toolbox's files which makes extensions complicated and redundant (e.g. name of
model in the path, the function name, the plotting routine).

Thanks
David

[rgeirhos]

If you're looking to add many models to the toolbox, the following approach might work (the example is based on a resnet18 at different epochs, but could easily be adapted to other settings):

for epoch in range(0, 90):
    exec(f"""@register_model("pytorch")
def resnet18_epoch_{epoch}(model_name="resnet18", *args):
    model = torch.load(f"/path/to/model/model_name_epoch_{epoch}")
    model.load_state_dict(checkpoint)
    return PytorchModel(model, model_name, *args)
    """)

This snippet can be inserted in modelvshuman/models/pytorch/model_zoo.py and would be a concise version of registering 90 different models. Is this in line with what you had in mind?

[dekuenstle]

Thanks, this is a work-around that can simplify what we are doing.

My initial request was more towards a function that is available when importing the toolbox.
If you modify the example with a custom model, I would expect the following user experience:

import torch
from modelvshuman import Plot, Evaluate, register_model
from modelvshuman import constants as c
from plotting_definition import plotting_definition_template


def run_evaluation():
    register_model('pytorch', 'crazy-model', torch.load('path/to/model.pth'))
    models = ["resnet50", "bagnet33", "simclr_resnet50x1", 'crazy-model']
    datasets = c.DEFAULT_DATASETS # or e.g. ["cue-conflict", "uniform-noise"]
    params = {"batch_size": 64, "print_predictions": True, "num_workers": 20}
    Evaluate()(models, datasets, **params)