nailarais1/image-classifier-efficientnet

EfficientNet-B0 Model for Image Classification

This repository contains an EfficientNet-B0 model trained on a custom dataset for image classification tasks.

Model Details

• Architecture: EfficientNet-B0
• Input Size: 224x224 RGB images
• Number of Classes: 10
• Dataset: Custom dataset with 10 categories
• Optimizer: AdamW
• Loss Function: CrossEntropyLoss
• Validation Accuracy: 85.3%
• Device Used for Training: CUDA (GPU)

Usage

Load the Model To load the model, use the following code:

import torch

Load model and metadata
model = torch.load("efficientnet-results-and-model.pth", map_location="cpu")

Access class-to-index mapping

class_to_idx = model['class_to_idx']

Load the state dictionary
state_dict = model['model_state_dict']

Reconstruct EfficientNet-B0
from torchvision.models import efficientnet_b0
model = efficientnet_b0(weights=None)
model.classifier[1] = torch.nn.Linear(model.classifier[1].in_features, len(class_to_idx))
model.load_state_dict(state_dict)
model.eval()

print("Model successfully loaded!")

Training Details
Learning Rate: 0.001
Batch Size: 32
Epochs: 3
Augmentations:
Random Resized Crop
Horizontal Flip
Color Jitter
Normalization (mean: [0.485, 0.456, 0.406], std: [0.229, 0.224, 0.225])

Files in this Repository
best_model.pth: Trained model weights
efficientnet.json: Model configuration file
README.md: Documentation for this model
efficientnet.txt: Training Results

Acknowledgments
Framework: PyTorch
Pretrained Weights: TorchVision
Training: Mixed precision using torch.cuda.amp for efficient training on GPU.