---
tags:
- image-classification
- timm
- pytorch
- safetensors
library_name: timm
datasets:
- parlange/dark-energy-survey-supernova
base_model:
- timm/cait_s24_224.fb_dist_in1k
---

# CaiT Model (cait_s24_224)

This repository contains a fine-tuned `CaiT` model from the `timm` library, intended for binary image classification.

The model weights are available in both standard PyTorch (`.bin`) and SafeTensors (`.safetensors`) formats.

## Model Details

* **Architecture**: `cait_s24_224`
* **Original Library**: `timm`
* **Fine-tuning Task**: Binary Image Classification
* **Number of Classes**: 2

---
## Training Hyperparameters

The model was trained with the following settings:

| Hyperparameter            | Value                                     |
|:--------------------------|:------------------------------------------|
| Optimizer                 | `AdamW`  |
| Learning Rate Schedule    | `1e-4 with CosineLRScheduler` |
| Batch Size                | `128` |
| Total Epochs              | `20` |
| Early Stopping Patience   | `7` on validation loss |
| Loss Function             | CrossEntropyLoss w/ Label Smoothing (`0.1`) |

---
## Training Results

Here are the key **test** metrics for this model:

* **Test Accuracy**: 0.988
* **Test AUC**: 0.993
* **Test F1 Score**: 0.988
* **Best Epoch**: 12.000

---
## How to use with `timm`

You can load this model directly from the Hugging Face Hub using `timm.create_model`. The `config.json` in this repo provides all necessary metadata.

```python
import torch
import timm

# Ensure you have timm and huggingface_hub installed:
# pip install timm "huggingface_hub>=0.23.0"

# Load the model directly from the Hub
# The `pretrained=True` flag will download the weights and config automatically.
model = timm.create_model(
    'hf-hub:parlange/cait-autoscan',
    pretrained=True
)
model.eval()

# The model's default_cfg will now be populated with mean/std and input size
print(model.default_cfg)

# Example inference with a dummy input
dummy_input = torch.randn(1, 3, model.default_cfg['input_size'][-2], model.default_cfg['input_size'][-1])
with torch.no_grad():
    output = model(dummy_input)

print(f"Output shape: {output.shape}") # Should be torch.Size([1, 2])
print(f"Predictions: {torch.softmax(output, dim=1)}")
```

---
## Original Checkpoint

The original `.pth` checkpoint file used for this model is also available in this repository.