binary_classification_ontorch.py

# -*- coding: utf-8 -*-
"""binary_classification_OnTorch.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1sx9YkjK0us8xPucnbhOmaY5kGVxaHoj9

In this work, we perform fruit and vegetable image classification using PyTorch. We train and test several different models, measure inference time, and compare the results.

install the necessary packages and libraries
"""

from google.colab import drive
drive.mount('/content/drive')

!pip install pytorch_lightning
!pip install torchsummary
!pip install timm

!pip install torch torchvision albumentations

import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from pathlib import Path
import pandas as pd
import albumentations as A
from albumentations.pytorch import ToTensorV2
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import cv2
import pytorch_lightning as pl
import torchsummary as summary
from torchmetrics import functional as F_metrics
from pytorch_lightning.callbacks import BackboneFinetuning, EarlyStopping, LearningRateMonitor
from torch.optim.lr_scheduler import ReduceLROnPlateau
from sklearn.model_selection import train_test_split
from pytorch_lightning.loggers import CSVLogger
import timm
from pathlib import Path
import torchvision
from PIL import Image

torch.cuda.is_available()

ROOT = Path.cwd()
DATASET = ROOT.joinpath('drive/MyDrive/dataset/FruitsVegetables')
MODEL_PATH = ROOT.joinpath('drive/MyDrive/Models/torch_model')
MODEL_PATH.mkdir(parents=True, exist_ok=True)

files = [(img_file, img_file.parent.stem) for img_file in DATASET.glob('*/*.*')]
df = pd.DataFrame(
    data=files,
    columns=['ImgFile', 'Type']
)
df['Type_i'] = df.Type.astype('category').cat.codes
df.Type.value_counts().plot(kind='bar')
labels = df.groupby('Type_i')['Type'].first()
labels

"""As we can see, the classes are balanced and therefore we can continue."""

train, val_test = train_test_split(df, test_size=0.3, stratify=df.Type, random_state=43)
val, test = train_test_split(val_test, test_size=0.5, stratify=val_test.Type, random_state=43)

print(f"Train size = {train.shape[0]}, Validation size = {val.shape[0]}, Test size = {test.shape[0]}")

train.head()

class CustomDataset(Dataset):
    def __init__(self, df: pd.DataFrame):

        self.df = df
        self.base_transform = A.Compose(
                [

                A.SmallestMaxSize(224, p=1),
                A.CenterCrop(224,224, p=1),
                ToTensorV2(),
                ]
            )

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):

        img_path = self.df.iat[idx, 0]

        image = cv2.imread(img_path.as_posix())
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        label = self.df.iat[idx, 2]
        image = self.base_transform(image=image)['image']
        return image, label


train_dataset = CustomDataset(df=train)
val_dataset = CustomDataset(df=val,)
test_dataset = CustomDataset (df=test)

train_loader = DataLoader(train_dataset, batch_size=20, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=20, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=20, shuffle=False)

# Visualization example We display the first 4 images from the batch (imgs, lbls).

imgs, lbls = next(iter(train_loader))
plt.figure(figsize=(10, 10))
for i, (img, lbl) in enumerate(zip(imgs, lbls)):
    ax = plt.subplot(2, 2, i + 1)
    plt.imshow(transforms.ToPILImage()(img))
    ax.set_title(labels[lbl.numpy()])
    if i >= 3:
        break
plt.show()

device = 'cuda' if torch.cuda.is_available() else 'cpu'
device

"""The code below
prepares two transformation schemes (transforms) for images. The first (basic_transforms) is without augmentation, the second (aug_transforms) is with augmentation (horizontal flip and random rotation).

basic_transforms is suitable for basic data loading (without distortion), and aug_transforms adds random transformations to increase the diversity of the training set.

"""

IMAGE_SIZE = (224, 224)
BATCH_SIZE = 32

basic_transforms = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
   ])


aug_transforms = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.RandomHorizontalFlip(p=0.5),
    transforms.RandomRotation(degrees=15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
   ])

test.to_csv(DATASET.joinpath('test.csv'), index=False)
print(f"train.shape={train.shape}, val.shape={val.shape}, test.shape={test.shape}")

test_df = pd.read_csv(DATASET.joinpath('test.csv'))

import pandas as pd


train.to_csv(DATASET.joinpath('train.csv'), index=False)
val.to_csv(DATASET.joinpath('val.csv'), index=False)

print("Названия столбцов в df_train:", df_train.columns)

df_train = pd.read_csv(DATASET.joinpath('train.csv'))
df_val = pd.read_csv(DATASET.joinpath('val.csv'))

unique_labels = df_train['Type_i'].unique()

"""The code for image transformation was based on the official documentation of the torchvision library, which describes the basic transformations and their application in CV tasks. Here is the link.
- https://pytorch.org/vision/stable/transforms.html

And about PyTorch
- https://www.nvidia.com/en-us/glossary/pytorch/
- https://www.ibm.com/think/topics/pytorch

code below
Reads a CSV file containing information about image paths and class labels. Then prints the unique values ​​of these labels and the first few rows of the “Label_i” column values.

This line loads data from a CSV file that contains image paths and their class labels. This allows you to work with data in the convenient DataFrame format that Pandas provides.

Next, we load the data. There is a class called Dataset that we inherit from. In this class, we must implement two methods:

- the magic method len that returns the length of our dataset
- getitem getitem : a method that allows you to get elements by their index (method for the iterator)
"""

class FruitsVegetables(Dataset):
    def __init__(self, df: pd.DataFrame, transform=None):

        self.df = df
        self.transform = transform

    def __len__(self):
        return len(self.df)

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        img_path = row["ImgFile"]
        label = row["Type_i"]


        image = Image.open(img_path).convert('RGB')

        if self.transform is not None:
            image = self.transform(image)

        return image, label

"""The `FruitsVegetables` class is based on the descriptions in the official PyTorch docs about custom datasets and data loaders
-  https://pytorch.org/docs/stable/data.html#torch.utils.data.Dataset
- https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader
- https://pytorch.org/tutorials/beginner/basics/data_tutorial.html

then we create a dataloader - this is a batch former. What it does is it takes our dataset that we made (in this case, Dataset) and forms it into batches.
"""

train_dataset = FruitsVegetables(df_train, transform=basic_transforms)
val_dataset = FruitsVegetables(df_val, transform=basic_transforms)
test_dataset = FruitsVegetables(test_df, transform=basic_transforms)

train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)

# for augmentation
train_aug_dataset = FruitsVegetables(df_train, transform=aug_transforms)
train_aug_loader = DataLoader(train_aug_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)

"""PyTorch docs on data loading  https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader

### 2 Defining models

Below I create a custom model that consists of several convolutional layers and a fully connected layer. This model does not use pre-trained layers, but instead relies on its own constructed layers.
"""

# 2.1. homemade (custom) model without augmentation

class CustomModel(nn.Module):

    def __init__(self):

        super().__init__()

        self.conv1 = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        self.fc1 = nn.Linear(64 * 56 * 56, 128)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):

        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        x = self.pool(x)
        x = torch.flatten(x, start_dim=1)
        x = F.relu(self.fc1(x))

        x = self.fc2(x).squeeze()

        return x

"""This BELOW model is more complex and uses additional layers to extract features, improving performance on more complex data and applying regularization mechanisms to combat overfitting.

"""

# 2.2. homemade (custom) model with augmentation

class AugmentedCustomModel(nn.Module):

    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2d(2)  # 224 -> 112

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2d(2)  # 112 -> 56

        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.pool3 = nn.MaxPool2d(2)  # 56 -> 28

        self.conv5 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        # 256 x 28 x 28

        self.fc1 = nn.Linear(256 * 28 * 28, 128)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(128, 1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)

        x = F.relu(self.conv2(x))
        x = self.pool2(x)

        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.pool3(x)

        x = F.relu(self.conv5(x))
        x = torch.flatten(x, start_dim=1)

        x = F.relu(self.fc1(x))
        x = self.dropout(x)

        x = self.fc2(x).squeeze()  # logits
        return x

"""In binary classification models, we remove the sigmoid at the end because we use the BCEWithLogitsLoss loss function. This function itself applies the sigmoid to the output (logits) and calculates the loss. Thus, we do not need to use the sigmoid in the model, which simplifies the calculations and can make the training more stable. This saves time and resources when training the network.

The code for the custom models was developed based on the PyTorch documentation on the `torch.nn.Module` class, which describes how to create and use custom neural networks here https://pytorch.org/docs/stable/generated/torch.nn.Module.html

Next we create a class `VGG16Model` that inherits from `nn.Module`.
"""

# 2.3. vgg16 v2

class VGG16Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.base_model = models.vgg16(weights='DEFAULT')

        self.base_model.classifier[6] = nn.Linear(4096, 1)

    def forward(self, x):

        x = self.base_model(x)

        return x.squeeze()

"""- information about pre-training models available, documentation about vgg16 https://pytorch.org/vision/stable/models.html#torchvision.models.vgg16

Next - The `TransferLearnModel` class is created, inheriting from `nn.Module`.
- We use MobileNetV3 as the base model and replace the last layer to fit the binary classification task.
- The `forward` method also defines how the data flows through the model, returning logits.
"""

# 2.4. TransferLearnModel

class TransferLearnModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.base_model = models.mobilenet_v3_large(weights='DEFAULT')

        self.base_model.classifier[-1] = nn.Linear(1280, 1)

    def forward(self, x):

        x = self.base_model(x)
        return x.squeeze()

""" Here we create an `EfficientNetModel` class that uses the `timm` library to load the EfficientNet model.
- We get the input feature count of the last layer and replace it with a new layer for binary classification.
- The `forward` method determines how the data flows through the model, returning logits.
"""

class EfficientNetModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.base_model = timm.create_model('efficientnet_b0', pretrained=True)

        out_features = self.base_model.classifier.in_features

        self.base_model.classifier = nn.Linear(out_features, 1)

    def forward(self, x):
        x = self.base_model(x)

        return x.squeeze()

"""The `MobileNetV3Model` class is created, which uses the pre-trained MobileNetV3 architecture.
- The last layer is replaced with a new one, which corresponds to the binary classification task.
- The `forward` method gets the logits from the base model and returns them.
"""

# 2.6 MobileNetV3Large v2

class MobileNetV3Model(nn.Module):

    def __init__(self):

        super().__init__()

        self.base_model = models.mobilenet_v3_large(weights="DEFAULT")

        in_features = self.base_model.classifier[-1].in_features

        self.base_model.classifier[-1] = nn.Linear(in_features, 1)

    def forward(self, x):

        x = self.base_model(x)
        return x.squeeze()

"""- documentation https://pytorch.org/vision/stable/models.html#torchvision.models.mobilenet_v3_large

### 3 Training
"""

# Commented out IPython magic to ensure Python compatibility.
# %%bash
# mkdir -p /content/drive/MyDrive/models

import os
os.makedirs('/content/drive/MyDrive/Models', exist_ok=True)

"""This function is designed to train a model using training and validation datasets."""

def train_model(  # v2
    model,
    model_name,
    train_loader,
    val_loader,
    epochs=10,
    lr=1e-4
):

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    criterion = nn.BCEWithLogitsLoss()

    optimizer = torch.optim.Adam(model.parameters(), lr=lr)

    best_val_loss = float('inf')
    best_val_acc = 0.0
    history = []

    print(f"\nTraining model: {model_name}")

    for epoch in range(1, epochs + 1):

        model.train()

        running_loss = 0.0

        for imgs, labels in train_loader:
            imgs, labels = imgs.to(device), labels.float().to(device)

            optimizer.zero_grad()
            logits = model(imgs)
            loss = criterion(logits, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()



        model.eval()
        val_loss = 0.0
        correct = 0
        total = 0

        with torch.no_grad():

            for imgs, labels in val_loader:
                imgs, labels = imgs.to(device), labels.float().to(device)

                logits = model(imgs)
                loss = criterion(logits, labels)
                val_loss += loss.item()

                probs = torch.sigmoid(logits)
                preds = (probs >= 0.5).float()
                correct += (preds == labels).sum().item()
                total += labels.size(0)


        avg_train_loss = running_loss / len(train_loader)
        avg_val_loss = val_loss / len(val_loader)
        val_acc = correct / total

        history.append({
            'epoch': epoch,
            'train_loss': avg_train_loss,
            'val_loss': avg_val_loss,
            'val_acc': val_acc
        })

        print(
            f"epoch [{epoch}/{epochs}]: "
            f"train_loss={avg_train_loss:.4f}, "
            f"val_loss={avg_val_loss:.4f}, "
            f"val_acc={val_acc:.4f}"
        )


        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            torch.save(model.state_dict(), f"{model_name}_best_loss.pth")

        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), f"{model_name}_best_acc.pth")

    print(
        f"Best val_loss for {model_name}: {best_val_loss:.4f}, "
        f"Best  val_acc: {best_val_acc:.4f}"
    )
    return history

"""BELOW
The code trains six different models (two custom and four pre-trained), storing their best results.

- Creates a DataFrame `df_results` to record the training results for all models.
- For each model:
- Initializes the model.
- Calls `train_model` to train the model and get a history of loss and accuracy.
- Extracts the best loss and accuracy from this history and writes them to the DataFrame.

- About`BCEWithLogitsLoss` https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html

"""

df_results = pd.DataFrame(columns=['Model_name', 'Best_loss', 'Best_acc'])

# custom wo augm
custom_model = CustomModel()
hist_custom = train_model(custom_model, "CustomModel_wo_augm", train_loader, val_loader, epochs=10)
best_loss_custom = min(h['val_loss'] for h in hist_custom)
best_acc_custom = max(h['val_acc'] for h in hist_custom)
df_results.loc[len(df_results)] = ['CustomModel_wo_augm', best_loss_custom, best_acc_custom]

#  custom w  augm
augm_model = AugmentedCustomModel()
hist_augm = train_model(augm_model, "CustomModel_w_augm", train_aug_loader, val_loader, epochs=10, lr=1e-4)
best_loss_augm = min(h['val_loss'] for h in hist_augm)
best_acc_augm = max(h['val_acc'] for h in hist_augm)
df_results.loc[len(df_results)] = ['CustomModel_w_augm', best_loss_augm, best_acc_augm]

# VGG16
vgg_model = VGG16Model()
hist_vgg = train_model(vgg_model, "VGG16Model", train_loader, val_loader, epochs=15, lr=1e-4)
best_loss_vgg = min(h['val_loss'] for h in hist_vgg)
best_acc_vgg = max(h['val_acc'] for h in hist_vgg)
df_results.loc[len(df_results)] = ['VGG16Model', best_loss_vgg, best_acc_vgg]

# Transfer learning
transfer_model = TransferLearnModel()
hist_transfer = train_model(transfer_model, "TransferLearnModel", train_loader, val_loader, epochs=10)
best_loss_transfer = min(h['val_loss'] for h in hist_transfer)
best_acc_transfer = max(h['val_acc'] for h in hist_transfer)
df_results.loc[len(df_results)] = ['TransferLearnModel', best_loss_transfer, best_acc_transfer]

# EfficientNetB0
eff_model = EfficientNetModel()
hist_eff = train_model(eff_model, "EfficientNetModel", train_loader, val_loader, epochs=10)
best_loss_eff = min(h['val_loss'] for h in hist_eff)
best_acc_eff = max(h['val_acc'] for h in hist_eff)
df_results.loc[len(df_results)] = ['EfficientNetModel', best_loss_eff, best_acc_eff]

# MobileNetV3Large
mobnet_model = MobileNetV3Model()
hist_mobnet = train_model(mobnet_model, "MobileNetV3Model", train_loader, val_loader, epochs=10)
best_loss_mobnet = min(h['val_loss'] for h in hist_mobnet)
best_acc_mobnet = max(h['val_acc'] for h in hist_mobnet)
df_results.loc[len(df_results)] = ['MobileNetV3Model', best_loss_mobnet, best_acc_mobnet]

print(df_results.index)

print(df_results)

"""I liked the article in terms of writing code and general understanding.
- https://machinelearningmastery.com/pytorch-tutorial-develop-deep-learning-models/

Now let's save the models. I chose the ONNX format. Why? If you want your model to be able to run outside of PyTorch, for example in TensorFlow or Caffe2, ONNX allows you to do this. ONNX models can be optimized using tools like ONNX Runtime, which allows you to increase performance during inference. But if you plan to continue training the model or perform additional training, it is better to save it in the PyTorch format.
"""

!pip install onnx

# Commented out IPython magic to ensure Python compatibility.
import onnx
import torch
import torchvision.models as tv_models

# %load_ext autoreload
# %autoreload 2

from torch.onnx import export


model_classes = {
    'CustomModel_wo_augm': CustomModel,
    'CustomModel_w_augm': AugmentedCustomModel,
}


MODEL_PATH = Path("/content/drive/MyDrive/Models/torch_model")


for name, ModelCls in model_classes.items():
    model = ModelCls()

    ckpt_file = MODEL_PATH / f"{name}_best_acc.pth"
    if ckpt_file.exists():
        model.load_state_dict(torch.load(ckpt_file, map_location='cpu'),
                              strict=False)
        print(f"weights loaded for {name}")
    else:
        print(f"NO weights for {name}, exporting random-init model")

    model.eval()


    export(  # onnx.export
        model,
        torch.randn(1, 3, 224, 224),           # dummy-int
        (MODEL_PATH / f'{name}.onnx').as_posix(),
        opset_version=14,
        input_names = ['input'],
        output_names= ['output'],
        dynamic_axes={'input': {0: 'batch'},
                      'output': {0: 'batch'}}
    )

device = torch.device('cuda')
vgg_model.to(device).eval()

dummy_input = torch.randn(1, 3, 224, 224, device=device)

torch.onnx.export(
    vgg_model,
    dummy_input,
    (MODEL_PATH / "VGG16Model.onnx").as_posix(),
    opset_version=14,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input':  {0: 'batch'},
                  'output': {0: 'batch'}}
)

device = torch.device('cuda')
transfer_model.to(device).eval()

dummy_input = torch.randn(1, 3, 224, 224, device=device)

torch.onnx.export(
    transfer_model,
    dummy_input,
    (MODEL_PATH / "TransferLearnModel.onnx").as_posix(),
    opset_version=14,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input':  {0: 'batch'},
                  'output': {0: 'batch'}}
)

device = torch.device('cuda')
eff_model.to(device).eval()

dummy_input = torch.randn(1, 3, 224, 224, device=device)

torch.onnx.export(
    eff_model,
    dummy_input,
    (MODEL_PATH / "EfficientNetModel.onnx").as_posix(),
    opset_version=14,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input':  {0: 'batch'},
                  'output': {0: 'batch'}}
)

device = torch.device('cuda')
mobnet_model.to(device).eval()

dummy_input = torch.randn(1, 3, 224, 224, device=device)

torch.onnx.export(
    mobnet_model,
    dummy_input,
    (MODEL_PATH / "MobileNetV3Model.onnx").as_posix(),
    opset_version=14,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input':  {0: 'batch'},
                  'output': {0: 'batch'}}
)

df_results

"""- Oficial documentation on loading save models, state https://pytorch.org/tutorials/beginner/saving_loading_models.html
- examples https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html
- in ONNX https://pytorch.org/tutorials/advanced/super_resolution_with_onnxruntime.html

### 5 Testing

This function evaluates the performance of the passed model on the test data and returns the loss and accuracy.
- Parameters:
- `model`: the trained model to test.
- `model_name`: the name of the model to output.
- `test_loader`: the test data loader.
"""

from pathlib import Path
import torch
import numpy as np

def test_model_onnx(onnx_path: Path, loader):

    sess = ort.InferenceSession(onnx_path.as_posix(),
                                providers=['CPUExecutionProvider'])
    in_name  = sess.get_inputs()[0].name
    out_name = sess.get_outputs()[0].name

    correct, total = 0, 0
    for imgs, labels in loader:
        logits = sess.run(
            [out_name],
            {in_name: imgs.numpy().astype(np.float32)}
        )[0]
        preds   = (torch.sigmoid(torch.from_numpy(logits)) >= 0.5).numpy()
        correct += (preds == labels.numpy()).sum()
        total   += labels.size(0)

    return correct / total

"""### 6 Function for measuring inference time for one image"""

import time
import torch

def measure_inference_time(model, device, loader): ###
    model.eval().to(device)
    dummy = next(iter(loader))[0][0:1].to(device)

    with torch.no_grad():
        for _ in range(5):
            _ = model(dummy)

    if device.type == 'cuda':
        torch.cuda.synchronize()
    t0 = time.time()
    with torch.no_grad():
        _ = model(dummy)
    if device.type == 'cuda':
        torch.cuda.synchronize()
    return time.time() - t0

"""### 7 Testing models and adding results to df_results"""

!pip install --quiet onnxruntime   # CPU-vers
# pip install --quiet onnxruntime-gpu  # еGPU-vers

model_classes = {
    'CustomModel_wo_augm': CustomModel,
    'CustomModel_w_augm': AugmentedCustomModel,
    'VGG16Model'        : VGG16Model,
    'TransferLearnModel': TransferLearnModel,
    'EfficientNetModel' : EfficientNetModel,
    'MobileNetV3Model'  : MobileNetV3Model,
}

import onnxruntime as ort
import time, torch, pandas as pd

def test_model_onnx(onnx_path: Path, loader):
    sess = ort.InferenceSession(onnx_path.as_posix(),
                                providers=['CPUExecutionProvider'])
    in_name  = sess.get_inputs()[0].name
    out_name = sess.get_outputs()[0].name
    correct, total = 0, 0
    for imgs, labels in loader:
        logits = sess.run([out_name],
                          {in_name: imgs.numpy().astype('float32')})[0]
        preds = (torch.sigmoid(torch.from_numpy(logits)) >= 0.5).numpy()
        correct += (preds == labels.numpy()).sum()
        total   += labels.size(0)
    return correct / total

def measure_inference_time(model, device, loader):
    model.to(device).eval()
    dummy = next(iter(loader))[0][0:1].to(device)
    with torch.no_grad():
        for _ in range(5): _ = model(dummy)
    if device.type == 'cuda': torch.cuda.synchronize()
    t0 = time.time()
    with torch.no_grad(): _ = model(dummy)
    if device.type == 'cuda': torch.cuda.synchronize()
    return time.time() - t0

df_results = pd.DataFrame({
    'Model_name': list(model_classes.keys()),
    'Test_acc'  : None,
    'Inference_time_1img': None
})

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

for idx, name in enumerate(df_results['Model_name']):
    onnx_file = MODEL_PATH / f'{name}.onnx'
    if onnx_file.exists():
        df_results.loc[idx, 'Test_acc'] = test_model_onnx(onnx_file, test_loader)
    else:
        print(f'ONNX file for {name} not found')
        continue

    torch_model = model_classes[name]()
    ckpt = MODEL_PATH / f'{name}_best_acc.pth'
    if ckpt.exists():
        torch_model.load_state_dict(torch.load(ckpt, map_location=device),
                                    strict=False)
    df_results.loc[idx, 'Inference_time_1img'] = \
        measure_inference_time(torch_model, device, test_loader)

print(df_results)
df_results.to_csv('pytorch_results.csv', index=False)

"""An article about testing models using transfer learning, including loading pre-trained weights and methods for evaluating models on test data, helped with the code. https://pytorch.org/tutorials/beginner/transfer_learning_tutorial.html"""

df_results

"""In total, EfficientNetModel showed the highest accuracy (about 90.42%), but when choosing the optimal model for a more realistic task, one should take into account not only the quality of classification, but also the speed of inference, as well as available resources.

The EfficientNetB0 model seems to be more flexible and suitable in terms of time and quality, while simple home-made networks are suitable for extremely tight constraints on model size or resources, but without claims to maximum accuracy.

In terms of inference time, the fastest are the "light" homemade models (around 0.001–0.002 s), but they are inferior in accuracy to deeper architectures due to a smaller number of layers and the absence of pre-trained weights. The MobileNetV3 and EfficientNetB0 network architectures, oriented towards mobile devices, also showed high speed (~0.008–0.011 s) with excellent quality.

"""