File size: 1,640 Bytes
df8e564 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
---
license: mit
language:
- en
base_model:
- google/efficientnet-b0
---
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. |