Model Card: MRI Brain Tumor Classification Model

Model Details

Architecture: EfficientNet-B1-based MRI classification model
Dataset: Brain Tumor MRI Dataset
Batch Size: 32
Loss Function: Triplet Margin Loss with Cosine Similarity
Optimizer: Adam (learning rate = 1e-2)

Model Architecture

This model is based on EfficientNet-B1 and has been modified for MRI brain tumor classification. The main adaptations include:

Modifications:

Input Channel Adjustment: The first convolutional layer is changed to accept single-channel (grayscale) MRI scans.
Classifier Head: The default classifier is replaced with a custom MLP featuring:
- Fully connected layers with 1280 → 756 → 256 units.
- SiLU activation.
- Batch normalization.
- Dropout for regularization.

Triplet Loss for Metric Learning:

The model uses Triplet Margin Loss with Cosine Similarity to learn an embedding space where MRI images of the same class are closer together, while images from different classes are farther apart.

Implementation

Model Definition

import torch
import torch.nn as nn
from torchvision.models import efficientnet_b1
from torch.nn import TripletMarginWithDistanceLoss
from torch.nn.functional import cosine_similarity

class MRIModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super(MRIModel, self).__init__()
        self.base_model = efficientnet_b1(weights=False)
        self.base_model.features[0] = nn.Sequential(
            nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), bias=False),
            nn.BatchNorm2d(32),
            nn.ReLU6(inplace=True),
        )
        self.base_model.classifier = nn.Sequential(
            nn.Linear(1280, 756),
            nn.SiLU(),
            nn.BatchNorm1d(756),
            nn.Dropout(0.2),
            nn.Linear(756, 256),
        )

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

Training Configuration

Batch Size: 32
Loss Function: Triplet Margin Loss (Cosine Similarity)
Optimizer: Adam (learning rate = 1e-2)

This model has been pushed to the Hub using the PytorchModelHubMixin integration: