# How to Fine-Tune RWKV-CLIP on Domain-Specific Datasets: Medical and Satellite Imaging Guide

> Learn how to fine-tune RWKV-CLIP on specialized datasets like medical and satellite images. Adapt the vision encoder and run distributed training for optimal results.

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

---

**Fine-tune RWKV-CLIP on domain-specific datasets like medical or satellite imagery by loading pretrained weights through `create_RWKV_Model`, adapting the vision encoder in [`model/Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Image_rwkv.py) to your resolution, and running the distributed training loop in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) with domain-optimized hyperparameters.**

RWKV-CLIP combines a **RWKV-based text encoder** with a **RWKV-driven vision encoder** to deliver efficient vision-language pretraining with recurrent-style inference speed. For practitioners working with **domain-specific datasets** such as medical X-rays or satellite photography, fine-tuning this architecture requires understanding its unique RWKV block structure and CUDA kernel dependencies. This guide walks through adapting the deepglint/rwkv-clip codebase for specialized imaging domains using the original contrastive training pipeline.

## Understanding the RWKV-CLIP Architecture for Domain Adaptation

### Text and Vision Encoders (`Text_RWKV` and `Image_RWKV`)

The model uses two parallel RWKV encoders defined in [`model/Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Text_rwkv.py) (lines 83-120) and [`model/Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Image_rwkv.py) (lines 53-80). The **text encoder** processes token embeddings through stacked RWKV blocks with time-mixing layers, while the **vision encoder** applies patch embeddings and learnable positional encodings before feeding patches through `Block_V6` instances. Both encoders output L2-normalized embeddings that feed into the contrastive loss.

### RWKV Blocks with Custom CUDA Kernels

Each encoder relies on **RWKV blocks** (`Block_V6`) containing **spatial-mix** (`VRWKV_SpatialMix_V6`) and **channel-mix** (`VRWKV_ChannelMix_V6`) sub-layers, as implemented in [`model/Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Image_rwkv.py) (lines 98-130). These blocks utilize the `RUN_CUDA_RWKV6` kernel (defined in `model/cuda_image/wkv6_cuda.cu` and `model/cuda_text/wkv6_cuda.cu`) for efficient recurrent computation during training and inference.

### Contrastive Loss Mechanism (`ClipLoss`)

The training objective resides in [`loss.py`](https://github.com/deepglint/rwkv-clip/blob/main/loss.py) (lines 1-80), where `ClipLoss` computes cosine similarity between image and text embeddings, scales them by a learnable **temperature parameter** (`logit_scale`), and supports distributed training optimizations like *gather-with-grad*.

## Preparing Domain-Specific Data for Fine-Tuning

### Building a Custom PyTorch Dataset

For medical or satellite data, implement a custom `Dataset` class that uses OpenCLIP's image transforms and CLIP's tokenizer:

```python
from torch.utils.data import Dataset, DataLoader
from PIL import Image
import clip
from open_clip.transform import image_transform

class MedImageDataset(Dataset):
    def __init__(self, img_paths, captions, transform):
        self.img_paths = img_paths
        self.captions = captions
        self.transform = transform
        self.tokenizer = clip.tokenize

    def __len__(self):
        return len(self.img_paths)

    def __getitem__(self, idx):
        img = Image.open(self.img_paths[idx]).convert('RGB')
        img = self.transform(img)
        txt = self.tokenizer([self.captions[idx]])[0]   # (77,)

        return img, txt

# Example usage

train_dataset = MedImageDataset(img_paths=med_img_list,
                               captions=med_cap_list,
                               transform=image_transform(cfg.input_size, True))
train_loader = DataLoader(train_dataset,
                         batch_size=cfg.batch_size,
                         shuffle=True,
                         num_workers=cfg.workers,
                         pin_memory=True)

```

### Configuring Image Resolution and Patch Size

Domain-specific imagery often requires resolution adjustments. The vision encoder's behavior is controlled via `model_config/*.json` files (e.g., [`RWKV_CLIP_B32.json`](https://github.com/deepglint/rwkv-clip/blob/main/RWKV_CLIP_B32.json)), which specify `input_size`, `image_patch_size`, and `n_embd`. For high-resolution satellite images, increase `input_size` to 384 or 448 and reduce `image_patch_size` to 16 for finer granularity.

## Fine-Tuning Implementation Steps

### Loading Pretrained Weights with `create_RWKV_Model`

Initialize the model using the utility function from [`model/__init__.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/__init__.py):

```python
import torch
from open_clip.transform import image_transform
import clip
from model.utils import create_RWKV_Model
from model_config.utils_notebook import load_model_configs

# Load config and model weights

cfg = load_model_configs('model_config/RWKV_CLIP_B32.json')
model = create_RWKV_Model(cfg, model_weight_path='model_31.pt').cuda()
model.eval()

# Preprocess an example image

transform = image_transform(cfg.input_size, False)
image = transform(Image.open('my_medical_image.png')).unsqueeze(0).cuda()

# Tokenize a domain-specific caption

text = clip.tokenize(["an MRI scan of the brain"]).cuda()

# Forward pass

with torch.no_grad():
    img_feat, txt_feat, logit_scale = model(image, text)
    img_feat = img_feat / img_feat.norm(dim=-1, keepdim=True)
    txt_feat = txt_feat / txt_feat.norm(dim=-1, keepdim=True)
    sim = (logit_scale * img_feat @ txt_feat.T).softmax(dim=-1)
print('Similarity:', sim.item())

```

### Modifying the Training Pipeline ([`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py))

The main training entry point in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) (lines 18-46) supports custom data loaders through the `--train-data` argument. You can point this to your custom dataset module or use the DALI pipeline structure shown in the dataloaders directory:

```bash

# Save the custom loader as a Python module, e.g. my_dataloaders.py

python train.py \
  --output ./finetuned_rwkv_clip \
  --train-data ./my_dataloaders.py:train_loader \
  --train-num-samples 20000 \
  --epochs 12 \
  --lr 0.01 \
  --batch-size 64 \
  --precision bf16 \
  --image-patch-size 16

```

### Hyperparameter Adjustments for Medical and Satellite Data

Adapt these key arguments in [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) or your config JSON:

- **`--input-size`**: Increase to 384+ for satellite imagery to preserve spatial details.
- **`--image-patch-size`**: Use 16 instead of 32 to capture fine-grained medical structures.
- **`--drop-path-rate`**: Raise to 0.4 for small medical datasets to prevent overfitting.
- **`--lr`**: Reduce to 0.001 for medical imaging where pretrained features need delicate adjustment.
- **`--with-cp`**: Enable gradient checkpointing to train larger models on limited GPU memory.
- **`--precision bf16`**: Maintain bfloat16 mixed precision for stable training on modern GPUs.

## Optimization Strategies for Specialized Domains

### Staged Freezing: Vision-First Training

To preserve linguistic knowledge while adapting visual features, freeze the **text encoder** initially by setting `requires_grad=False` on parameters loaded via `Text_RWKV`. Train only the vision encoder (`Image_RWKV`) for the first few epochs, then unfreeze both encoders for joint fine-tuning. This staged approach prevents catastrophic forgetting of general text representations when learning domain-specific visual patterns.

### Handling High-Resolution Imagery

Medical and satellite datasets often exceed standard 224×224 resolutions. Adjust the `input_size` parameter in your model config and ensure your custom dataset's transform uses `image_transform(cfg.input_size, is_train=True)`. The patch embedding layer in `Image_RWKV` automatically adapts to different input resolutions as long as the patch size remains consistent with the pretrained configuration.

## Summary

- **RWKV-CLIP** uses unified RWKV blocks for both text ([`model/Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Text_rwkv.py)) and vision ([`model/Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Image_rwkv.py)) encoders, enabling efficient fine-tuning with custom CUDA kernels.
- **Domain adaptation** requires adjusting `input_size`, `image_patch_size`, and learning rates in `model_config/*.json` or via command-line arguments to [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py).
- **Custom datasets** should implement OpenCLIP transforms and CLIP tokenization, then integrate with the training loop through `--train-data` specifications.
- **Staged freezing** of the text encoder during initial epochs stabilizes fine-tuning on small medical or satellite datasets.
- The **contrastive loss** in [`loss.py`](https://github.com/deepglint/rwkv-clip/blob/main/loss.py) and mixed-precision training support (`bf16`) allow scalable fine-tuning on domain-specific vision-language tasks.

## Frequently Asked Questions

### What makes RWKV-CLIP suitable for domain-specific fine-tuning?

RWKV-CLIP shares the same **RWKV block architecture** between its text and vision encoders, allowing both modalities to benefit from fast, recurrent-style inference while maintaining parallel training speed. This design, implemented in [`model/Image_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Image_rwkv.py) and [`model/Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Text_rwkv.py), enables end-to-end fine-tuning with a single optimizer loop, making it straightforward to adapt to medical or satellite imagery using the standard [`train.py`](https://github.com/deepglint/rwkv-clip/blob/main/train.py) pipeline.

### How do I adapt RWKV-CLIP for high-resolution satellite images?

Modify the `input_size` parameter in your model configuration JSON (e.g., [`model_config/RWKV_CLIP_B32.json`](https://github.com/deepglint/rwkv-clip/blob/main/model_config/RWKV_CLIP_B32.json)) to 384 or 448, and set `image_patch_size` to 16 for finer granularity. Update your dataset's transform to use `image_transform(cfg.input_size, is_train=True)` from `open_clip.transform`, and ensure your GPU memory can accommodate the larger sequence length or enable `--with-cp` for gradient checkpointing.

### Can I freeze the text encoder during initial fine-tuning?

Yes. After loading the model via `create_RWKV_Model` in [`model/__init__.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/__init__.py), iterate through the text encoder parameters (instantiated from `Text_RWKV` in [`model/Text_rwkv.py`](https://github.com/deepglint/rwkv-clip/blob/main/model/Text_rwkv.py)) and set `param.requires_grad = False` for the first few epochs. This preserves general linguistic knowledge while the vision encoder (`Image_RWKV`) adapts to domain-specific visual features, then unfreeze both for joint optimization.

### What CUDA kernels are required for training RWKV-CLIP?

The model requires the custom **RWKV6 CUDA kernels** (`RUN_CUDA_RWKV6`) defined in `model/cuda_image/wkv6_cuda.cu` for the vision encoder and `model/cuda_text/wkv6_cuda.cu` for the text encoder. These kernels power the spatial-mix and channel-mix operations within `Block_V6`, enabling efficient recurrence computation during both forward and backward passes.