File size: 1,698 Bytes
1a3a8ef 2e9036c 1a3a8ef | 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 | ---
pipeline_tag: zero-shot-image-classification
tags:
- open_clip
- clip
- vision
- image-text-retrieval
- laion400m
---
# ViT-B-32 OpenCLIP Model on LAION-400M
This is a ViT-B-32 model trained using [OpenCLIP](https://github.com/mlfoundations/open_clip) on the LAION-400M dataset.
## Training Details
The model was trained with the following configuration:
- **Model Architecture**: ViT-B-32
- **Dataset**: LAION-400M
- **Number of Samples**: 400M (~ 268,836,185 filtered samples used)
- **Hardware**: 2 Nodes, each with 4 H200 141GB GPUs (Total 8 GPUs)
- **Batch Size (per GPU)**: 4096
- **Precision**: `amp_bfloat16`
- **Total Epochs**: 32
- **Warmup Steps**: 2000
Additional specific performance-enhancing flags enabled during training: `--torchcompile`, `--local-loss`, and `--gather-with-grad`.
## Evaluation
- **Eval Epoch**: 0
- **imagenet-zeroshot-val-top1**: 0.6086
- **imagenet-zeroshot-val-top5**: 0.8632
## Usage
```python
import torch
import open_clip
# Load the model directly from huggingface
model, _, preprocess = open_clip.create_model_and_transforms('ViT-B-32', pretrained='hf-hub:lingkai/open-clip')
tokenizer = open_clip.get_tokenizer('ViT-B-32')
# Example inference
image = preprocess(Image.open("astronaut.png")).unsqueeze(0)
text = tokenizer(["a diagram", "a dog", "a cat"])
with torch.no_grad(), torch.cuda.amp.autocast():
image_features = model.encode_image(image)
text_features = model.encode_text(text)
image_features /= image_features.norm(dim=-1, keepdim=True)
text_features /= text_features.norm(dim=-1, keepdim=True)
text_probs = (100.0 * image_features @ text_features.T).softmax(dim=-1)
print("Label probs:", text_probs)
```
|