---
license: mit
tags:
  - image-classification
  - pytorch
  - resnet18
  - waste-classification
datasets:
  - trashnet
---

# TrashNet ResNet18

Fine-tuned **ResNet18** (ImageNet-pretrained backbone, custom classification head) for 6-class waste image classification on the [TrashNet](https://github.com/garythung/trashnet) dataset.

Used by the [Trash Classification System](https://github.com/yutongyu-ai/Trash_Classification_System) — a FastAPI + Streamlit app that serves this model for real-time inference.

## Classes

`cardboard`, `glass`, `metal`, `paper`, `plastic`, `trash`

## Training

- Backbone: `torchvision.models.resnet18` (IMAGENET1K_V1 weights), custom `Linear(64) -> ReLU -> Dropout(0.5) -> Linear(num_classes)` head
- Optimizer: AdamW + CosineAnnealingLR
- Class-weighted cross-entropy loss (TrashNet's `trash` class is underrepresented ~3.6x vs `paper`)
- Hyperparameters (`lr`, `weight_decay`, `batch_size`) selected via Optuna hyperparameter search (25 trials), then trained for 30 epochs with the winning config
- Trained on the University of Manchester CSF3 HPC cluster (SLURM, NVIDIA L40S GPU)

## Usage

```python
from huggingface_hub import hf_hub_download
import torch
from torchvision import models
import torch.nn as nn

def get_model(hidden_size=64, num_classes=6):
    model = models.resnet18(weights=None)
    model.fc = nn.Sequential(
        nn.Linear(model.fc.in_features, hidden_size),
        nn.ReLU(),
        nn.Dropout(p=0.5),
        nn.Linear(hidden_size, num_classes),
    )
    return model

checkpoint_path = hf_hub_download(
    repo_id="tonghahaha/trashnet-resnet18",
    filename="best_resnet18_trashnet.pth",
)
model = get_model(num_classes=6)
model.load_state_dict(torch.load(checkpoint_path, map_location="cpu", weights_only=True))
model.eval()
```