# How to Train RWKV-CLIP on Custom Image-Text Datasets Using the Official Training Scripts

> Learn to train RWKV-CLIP on custom image-text datasets. Convert data to RecordIO, configure DALI dataloader, and launch distributed training using the official scripts.

- Repository: [DeepGlint/rwkv-clip](https://github.com/deepglint/rwkv-clip)
- Tags: how-to-guide
- Published: 2026-02-28

---

**To train RWKV-CLIP on a custom dataset, convert your images and captions to MXNet RecordIO format using [`data2rec.py`](https://github.com/deepglint/rwkv-clip/blob/main/data2rec.py), configure the DALI dataloader in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py), and launch distributed training via [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) with your model configuration.**

The `deepglint/rwkv-clip` repository provides a complete pipeline for training vision-language models that use parallel RWKV backbones for image and text streams. By following the three core components—data preparation, GPU-accelerated loading, and distributed training—you can adapt the model to any custom image-text collection.

## Understanding the RWKV-CLIP Training Pipeline

Training RWKV-CLIP requires three integrated components that work together to process raw data and optimize the contrastive loss:

| Component | Purpose | Source File |
|-----------|---------|-------------|
| **Data preparation** | Converts raw images and captions into MXNet `.rec/.idx` files for efficient reading | [`data2rec.py`](https://github.com/deepglint/rwkv-clip/blob/main/data2rec.py) |
| **DALI dataloader** | Builds a GPU-accelerated pipeline that decodes JPEGs, applies augmentations, and returns image tensors with tokenized captions | [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) |
| **Training script** | Parses hyper-parameters, constructs the RWKV-CLIP model with image and text branches, wraps it with `DistributedDataParallel`, and runs the contrastive loop | [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) |

## Step 1: Prepare Your Custom Dataset

### Organize Images and Annotations

Store all images in a single directory (e.g., `mydata/images`). Create a JSON annotation file where each entry contains the image filename and a list of captions. The format follows the COCO loader expectations used in the repository:

```json
[
  {
    "image": "0001.jpg",
    "captions": ["A cat sitting on a mat", "A feline resting on carpet"]
  }
]

```

### Convert to MXNet RecordIO Format

The [`data2rec.py`](https://github.com/deepglint/rwkv-clip/blob/main/data2rec.py) script demonstrates how to pack images and tokenized captions into MXNet RecordIO files. For a production dataset, adapt the following pattern which concatenates multiple captions per image:

```python

# mydata/create_rec.py

import mxnet as mx
import cv2
import json
import os
import numpy as np
from src.open_alip import tokenize

idx = mx.recordio.MXIndexedRecordIO('mydata/data.idx', 'mydata/data.rec', 'w')

with open('mydata/annotations.json') as f:
    data = json.load(f)

for i, entry in enumerate(data):
    img = cv2.imread(os.path.join('mydata/images', entry['image']))
    
    # Tokenize all captions and flatten

    caps = [tokenize(c).flatten().numpy() for c in entry['captions']]
    labels = np.concatenate(caps)
    
    header = mx.recordio.IRHeader(flag=0, label=labels, id=1, id2=0)
    packed = mx.recordio.pack_img(header, img)
    idx.write_idx(i, packed)

```

Execute this script to generate `mydata/data.rec` and `mydata/data.idx`. The DALI dataloader in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) expects these files without the extension suffix.

## Step 2: Configure the Model and Data Pipeline

### Adjust Model Configuration

Copy an existing configuration from [`model_config/RWKV_CLIP_B32.json`](https://github.com/deepglint/rwkv-clip/blob/main/model_config/RWKV_CLIP_B32.json) and modify the architecture to match your dataset resolution and desired model capacity. Critical fields include:

- `input_size`: Image resolution (must match the DALI pipeline output size)
- `image_embed_dims` and `n_embd`: Shared hidden dimension for vision and language branches
- `image_patch_size`: Patch size for the vision RWKV encoder
- `n_layer`: Number of RWKV blocks in the text encoder

Example configuration snippet:

```json
{
    "input_size": 224,
    "image_embed_dims": 640,
    "image_patch_size": 32,
    "n_embd": 640,
    "n_layer": 6,
    "image_num_heads": 8
}

```

Save this as [`model_config/my_custom.json`](https://github.com/deepglint/rwkv-clip/blob/main/model_config/my_custom.json).

### Set Up the DALI Dataloader

The [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) file defines the `dali_dataloader` function, which constructs a GPU-accelerated pipeline using NVIDIA DALI. It reads the `.rec/.idx` files created in Step 1, decodes JPEGs on the GPU, applies random cropping and normalization, and returns image tensors paired with tokenized caption labels. Ensure your RecordIO files follow the format expected by the `mxnet_reader` ops defined in that file.

## Step 3: Launch Distributed Training

### Create the Training Script

The repository provides shell scripts in [`shell/train_RWKV_CLIP_B32_YFCC15M.sh`](https://github.com/deepglint/rwkv-clip/blob/main/shell/train_RWKV_CLIP_B32_YFCC15M.sh) that demonstrate the full torchrun command. Adapt the following template for your custom dataset, replacing paths and hyperparameters as needed:

```bash
#!/bin/bash

# train_my_custom.sh

# Model architecture parameters

input_size=224
image_depth=12
image_embed_dims=640
image_patch_size=32
image_hidden_rate=5
image_num_heads=8
drop_path_rate=0.3

# Text encoder parameters

ctx_len=77
vocab_size=49408
head_size=64
head_size_divisor=8
n_embd=640
n_layer=6
text_initialization=True

# Training parameters

ip_list=(127.0.0.1)          # Single-node training

lr=0.001
opt=adamw
weight_decay=0.2
train_num_samples=15000000   # Size of your dataset

epochs=32
batch_size=512
precision=bf16
open_checkpoint=False
traindata=mydata/data        # Path without .rec/.idx suffix

output=./output_my_custom
dropout=0.0

# Launch distributed training

torchrun --nproc_per_node 8 train.py \
    --drop-path-rate $drop_path_rate \
    --image-num-heads $image_num_heads \
    --input-size $input_size \
    --dropout $dropout \
    --precision $precision \
    --image-depth $image_depth \
    --image-embed-dims $image_embed_dims \
    --image-patch-size $image_patch_size \
    --image-hidden-rate $image_hidden_rate \
    --ctx-len $ctx_len \
    --vocab-size $vocab_size \
    --head-size $head_size \
    --head-size-divisor $head_size_divisor \
    --n-embd $n_embd \
    --n-layer $n_layer \
    --text-initialization $text_initialization \
    --batch-size $batch_size \
    --epochs $epochs \
    --lr $lr \
    --optimizer $opt \
    --output $output \
    --train-data $traindata \
    --train-num-samples $train_num_samples \
    --weight-decay $weight_decay &

```

Make the script executable and run it:

```bash
chmod +x train_my_custom.sh
./train_my_custom.sh

```

### Monitor Training Progress

The [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) script automatically handles logging and checkpointing:

- **Logs**: Written to `${output}/training.log` with loss values and learning rates.
- **TensorBoard**: Metrics are stored under `${output}/tensorboard` for visualization.
- **Checkpoints**: After each epoch, the script saves `RWKV_CLIP_model_<epoch>.pt` in the output directory (see the checkpointing logic in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py)).

View training metrics in real-time:

```bash
tensorboard --logdir ./output_my_custom/tensorboard

```

## Fine-Tuning vs. Training from Scratch

When launching training, you can choose between initializing from pretrained weights or training from scratch:

- **Fine-tuning**: Set `--open-checkpoint True` and ensure your `--output` directory contains a pretrained checkpoint file. The script will load existing weights before beginning the training loop.
- **From scratch**: Set `--open-checkpoint False` and the model will initialize with random weights according to the configuration in your JSON file.

For smaller custom datasets, keep regularization parameters modest:
- `--drop-path-rate`: Controls stochastic depth in the vision encoder (recommended 0.1-0.3 for fine-tuning).
- `--dropout`: Standard dropout rate for the text encoder (typically 0.0-0.1 for contrastive learning).

## Summary

- **Data preparation** is the first critical step: convert your images and captions to MXNet RecordIO format using the pattern shown in [`data2rec.py`](https://github.com/deepglint/rwkv-clip/blob/main/data2rec.py), ensuring each record contains tokenized caption arrays.
- **Configuration** requires copying [`model_config/RWKV_CLIP_B32.json`](https://github.com/deepglint/rwkv-clip/blob/main/model_config/RWKV_CLIP_B32.json) and adjusting `input_size`, `image_embed_dims`, and `n_embd` to match your hardware and dataset characteristics.
- **Training** is launched via `torchrun` using the template in [`shell/train_RWKV_CLIP_B32_YFCC15M.sh`](https://github.com/deepglint/rwkv-clip/blob/main/shell/train_RWKV_CLIP_B32_YFCC15M.sh), specifying your custom `--train-data` path (without file extensions) and monitoring via TensorBoard logs in the output directory.
- **Checkpointing** occurs automatically after each epoch, allowing you to resume training with `--open-checkpoint True` for fine-tuning scenarios.

## Frequently Asked Questions

### What format should my custom dataset be in for RWKV-CLIP?

RWKV-CLIP requires MXNet RecordIO format (`.rec` and `.idx` files) for efficient GPU loading via NVIDIA DALI. You must organize your raw images in a folder and create a JSON annotation file listing image filenames and their captions. Then use the pattern from [`data2rec.py`](https://github.com/deepglint/rwkv-clip/blob/main/data2rec.py) to tokenize captions with `src.open_alip.tokenize` and pack them into RecordIO records. The DALI dataloader in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) expects these files without the extension suffix in the `--train-data` argument.

### How do I resume training from a checkpoint in RWKV-CLIP?

To fine-tune or resume training, set the `--open-checkpoint True` flag in your training script and ensure the `--output` directory contains the previously saved checkpoint file (e.g., `RWKV_CLIP_model_<epoch>.pt`). The [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) script will load these weights before initializing the `DistributedDataParallel` wrapper. For training from scratch, set `--open-checkpoint False` and the model will initialize randomly according to your JSON configuration.

### What hardware requirements are needed to train RWKV-CLIP?

RWKV-CLIP training is optimized for multi-GPU setups using `torchrun` and NVIDIA DALI for GPU-accelerated data loading. The reference scripts assume 8 GPUs per node (`--nproc_per_node 8`) and use bfloat16 (`--precision bf16`) to reduce memory footprint. For the base configuration (`RWKV_CLIP_B32`), you need approximately 640-dimensional embeddings and 224px input resolution, requiring roughly 40-80GB of GPU memory depending on batch size (typically 512 per node). Single-node training is supported by setting `ip_list=(127.0.0.1)` in the shell script.

### How do I adjust the model architecture for different image resolutions?

Modify the `input_size` field in your configuration JSON (e.g., [`model_config/my_custom.json`](https://github.com/deepglint/rwkv-clip/blob/main/model_config/my_custom.json)) to match your desired resolution, ensuring it aligns with the `image_patch_size` so that `(input_size / image_patch_size)^2` yields an integer number of patches. For example, with `input_size: 224` and `image_patch_size: 32`, you get 49 patches. You must also update the DALI pipeline in [`dali.py`](https://github.com/deepglint/rwkv-clip/blob/main/dali.py) to output the same resolution, and adjust `image_embed_dims` and `n_embd` if you change the model capacity to handle different patch counts.