Evaluating MLflow

custom_predict_mnist_cnn.py

import torch
import torch.nn.functional as F
import torchvision
import torchvision.models as models
import mlflow
import mlflow.pytorch
import numpy as np


class Model(mlflow.pyfunc.PythonModel):
    def __init__(self):
        self.model = create_model()

    def predict(self, context, model_input):
        data = np.array(model_input.iloc[0][0]).reshape(1, 1, 28, 28)
        X = torch.Tensor(data)
        res = self.model(X).detach()
        return {'label': np.argmax(res.numpy())}


class Flatten(torch.nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)

 
def train(model, train_loader, loss_func, optimizer):
    model.train()
    step = 1
    for _ in range(0, 2):
        for batch_idx, (data, target) in enumerate(train_loader):
            optimizer.zero_grad()
            output = model(data)
            loss = loss_func(output, target)
            loss.backward()
            optimizer.step()
            print(f'{step}, Loss: {loss.item()}')
            mlflow.log_metric("train_loss", loss.item(), step=step)
            step += 1


def test(model, test_loader, loss_func):
    model.eval()
    correct = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            output = model(data)
            loss = loss_func(output, target)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            accuracy = (correct/len(test_loader.dataset)) * 100
            mlflow.log_metric("test_loss", loss.item(), step=batch_idx)
            mlflow.log_metric("accuracy", accuracy, step=batch_idx)

          
def create_model():
    inp = 32 * 7 * 7
    return torch.nn.Sequential(
        torch.nn.Conv2d(in_channels=1, out_channels=16,
                        kernel_size=5, stride=1, padding=2),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2),
        torch.nn.Conv2d(in_channels=16, out_channels=32,
                        kernel_size=5, stride=1, padding=2),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2),
        Flatten(),
        torch.nn.Linear(inp, 10))

 def main():
    m = Model()
    model = m.model
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    loss_func = torch.nn.CrossEntropyLoss()
    train_data = torchvision.datasets.MNIST(
        root='./mnist/',
        train=True,
        transform=torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize((0.1307,), (0.3081,))]),
        download=True,
    )
    test_data = torchvision.datasets.MNIST('./mnist/', train=False,
                                           transform=torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.1307,), (0.3081,))
    ]),
                               download=True)
    test_loader = torch.utils.data.DataLoader(test_data,
                                              batch_size=64)
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=64)
    cli = mlflow.tracking.MlflowClient()
    experiment = cli.get_experiment_by_name('mnist_cnn')
    with mlflow.start_run(experiment_id=experiment.experiment_id,
                          run_name='two_epoch_train_test') as run:
        train(model, train_loader, loss_func, optimizer)
        test(model, test_loader, loss_func)
        mlflow.set_tag("architecture", "CNN")
        mlflow.log_param("lr", 0.001)
        mlflow.log_param("batch_size", 64)
        run_id = run.info.run_id
        path = f"model/{run_id}"
        mlflow.pyfunc.save_model(python_model=m, path=path)


if __name__ == '__main__':
    main()

mnsit_ccn.py

import torch
import torch.nn.functional as F
import torchvision
import torchvision.models as models
import mlflow
import mlflow.pytorch


class Flatten(torch.nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)


def train(model, train_loader, loss_func, optimizer):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = loss_func(output, target)
        loss.backward()
        optimizer.step()
        print(f'{batch_idx}, Loss: {loss.item()}')
        mlflow.log_metric("train_loss", loss.item(), step=batch_idx)


def test(model, test_loader, loss_func):
    model.eval()
    correct = 0
    with torch.no_grad():
        for batch_idx, (data, target) in enumerate(test_loader):
            output = model(data)
            loss = loss_func(output, target)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()
            accuracy = (correct/len(test_loader.dataset)) * 100
            mlflow.log_metric("test_loss", loss.item(), step=batch_idx)
            mlflow.log_metric("accuracy", accuracy, step=batch_idx)


def main():
    inp = 32 * 7 * 7
    model = torch.nn.Sequential(
        torch.nn.Conv2d(in_channels=1, out_channels=16,
                        kernel_size=5, stride=1, padding=2),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2),
        torch.nn.Conv2d(in_channels=16, out_channels=32,
                        kernel_size=5, stride=1, padding=2),
        torch.nn.ReLU(),
        torch.nn.MaxPool2d(kernel_size=2),
        Flatten(),
        torch.nn.Linear(inp, 10))
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    loss_func = torch.nn.CrossEntropyLoss()
    train_data = torchvision.datasets.MNIST(
        root='./mnist/',
        train=True,
        transform=torchvision.transforms.Compose(
            [torchvision.transforms.ToTensor(),
             torchvision.transforms.Normalize((0.1307,), (0.3081,))]),
        download=True,
    )
    test_data = torchvision.datasets.MNIST('./mnist/', train=False,
                                           transform=torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize((0.1307,), (0.3081,))
    ]),
                               download=True)
    test_loader = torch.utils.data.DataLoader(test_data,
                                              batch_size=64)
    train_loader = torch.utils.data.DataLoader(train_data, batch_size=64)
    cli = mlflow.tracking.MlflowClient()
    experiment = cli.get_experiment_by_name('mnist_cnn')
    with mlflow.start_run(experiment_id=experiment.experiment_id,
                          run_name='one_epoch_train_test') as run:
        train(model, train_loader, loss_func, optimizer)
        test(model, test_loader, loss_func)
        mlflow.set_tag("architecture", "CNN")
        mlflow.log_param("lr", 0.001)
        mlflow.log_param("batch_size", 64)
        path = f"model/{run.info.run_id}"
        mlflow.pytorch.save_model(model, path)
        mlflow.log_artifact(f"{path}/data/model.pth")
        mlflow.log_artifact(f"{path}/MLmodel")


if __name__ == '__main__':
    main()