ML knowledge preservation for the AI era

# Create your dataloaders
train_dataloader, val_dataloader = your_function_to_get_dataloaders()

# Create a recommender object and fit to the training dataloader
recommender = modlee.recommender.from_modality_task(
    modality='image',
    task='classification',
    )

# Fit the recommender to the data meta-features
recommender.fit(train_dataloader)

# Get the model from the recommender and train
modlee_model = recommender.model
with modlee.start_run() as run:
    trainer = modlee.Trainer(max_epochs=10)
    trainer.fit(model=modlee_model,train_dataloaders=train_dataloader)

# Define your custom neural network
class ExampleCNN(torch.nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        self.conv1 = torch.nn.Conv2d(in_channels=3, out_channels=1,
                                     kernel_size=3, padding=1)
        self.fc1 = torch.nn.Linear(1 * 32 * 32, num_classes)
    def forward(self, x):
        x = torch.nn.functional.relu(self.conv1(x))
        x = x.view(-1, 1 * 32 * 32)
        x = self.fc1(x)
        return x

# Subclass the ModleeModel class to enable automatic documentation
class ModleeClassifier(modlee.model.ModleeModel):
    def __init__(self, model):
        super().__init__()
        self.model, self.loss_fn = model, torch.nn.functional.cross_entropy
    def forward(self, x):
        return self.model(x)
    def training_step(self, batch):
        x, y_target = batch
        y_pred = self(x)
        loss = self.loss_fn(y_pred, y_target)
        return {"loss": loss}
    def configure_optimizers(self):
        optimizer = torch.optim.SGD(self.parameters(), lr=0.001, momentum=0.9)
        return optimizer

# Create the model object
modlee_model = ModleeClassifier(ExampleCNN(10))

# Insert the modlee context manager before training
with modlee.start_run() as run:
    trainer = modlee.trainer(max_epochs=10)
    trainer.fit(model=modlee_model, train_dataloaders=train_dataloader)