lra-benchmark 0.1.2

A library for benchmarking lightweight Vision Transformers following the LRA methodology

Sure, I can provide you with a detailed README.md file and guide you through the steps to publish your library to PyPI.

LRA Benchmark for Image Classification on CIFAR-10

The LRA Benchmark for Image Classification on CIFAR-10 is a project that provides a benchmark for evaluating image classification models on the CIFAR-10 dataset. It includes a set of pre-defined models, as well as a framework for running benchmark experiments and reporting key metrics.

Project Structure

The project is structured as follows:

lra_benchmark/
├── lra_benchmark/
│   ├── models/
│   │   ├── base_model.py
│   │   └── tiny_vit.py
│   ├── dataset.py
│   ├── benchmark.py
│   ├── config.py
│   └── utils/
│       └── optimizer.py
├── examples/
│   └── benchmark_example.py
├── tests/
├── setup.py
└── README.md

• lra_benchmark/models/: Contains the implementation of the base model class (BaseModel) and the pre-defined model architectures, such as TinyViT.
• lra_benchmark/dataset.py: Provides functions for loading the CIFAR-10 dataset.
• lra_benchmark/benchmark.py: Defines the LRABenchmark class and the run_benchmark function, which are used to run the benchmark experiments.
• lra_benchmark/config.py: Defines the configuration options for the benchmark experiments.
• lra_benchmark/utils/optimizer.py: Provides utilities for creating optimizers.
• examples/benchmark_example.py: An example script that demonstrates how to use the library to run a benchmark experiment.
• tests/: Contains unit tests for the library.
• setup.py: The setup file for packaging and distributing the library.
• README.md: The project's documentation.

Using Your Own Models

To use your own model with the LRA Benchmark, you'll need to create a new model class that inherits from the BaseModel class. Here's an example:

from lra_benchmark.models.base_model import BaseModel
import torch.nn as nn

class MyModel(BaseModel):
    def __init__(self, img_size=32, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(32 * 8 * 8, num_classes)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

Once you've defined your custom model, you can use it in the same way as the TinyViT model in the benchmark_example.py script:

from lra_benchmark.models.my_model import MyModel
from lra_benchmark.config import LRABenchmarkConfig
from lra_benchmark.dataset import get_cifar10_datasets
from lra_benchmark.benchmark import run_benchmark

config = LRABenchmarkConfig(batch_size=128, max_epochs=5)
train_loader, val_loader, test_loader = get_cifar10_datasets(config.batch_size, config.num_workers)

model = MyModel(img_size=32, num_classes=10)
metrics = run_benchmark(config, model, train_loader, val_loader, test_loader)

print(f"Test Accuracy: {metrics['test_accuracy']:.2%}")
print(f"Test Loss: {metrics['test_loss']:.4f}")
print(f"Model Size: {metrics['model_size']} parameters")
print(f"Inference Time: {metrics['inference_time']:.4f} seconds")