Chinese-CLIP/
├── run_scripts/
│   ├── muge_finetune_vit-b-16_rbt-base.sh
│   ├── flickr30k_finetune_vit-b-16_rbt-base.sh
│   └── ...             # 更多finetune或评测脚本...
└── cn_clip/
    ├── clip/
    ├── eval/
    ├── preprocess/
    └── training/

${DATAPATH}
├── pretrained_weights/ # 存放对应模型ckpt
├── experiments/
├── deploy/	            # 用于存放ONNX & TensorRT部署模型
└── datasets/
    ├── MUGE/
    ├── Flickr30k-CN/
    └── .../            # 更多自定义数据集...

${DATAPATH}                     # 如：Flickr30k-CN/
└── datasets/
    └── ${dataset_name}/
        ├── train_imgs.tsv      # 图片id & 图片内容
        ├── train_texts.jsonl   # 文本id & 文本内容，连同匹配的图片id列表
        ├── valid_imgs.tsv
        ├── valid_texts.jsonl
        ├── test_imgs.tsv
        └── test_texts.jsonl

python cn_clip/preprocess/build_lmdb_dataset.py \
    --data_dir ${DATAPATH}/datasets/${dataset_name}
    --splits train,valid,test

${DATAPATH}
└── datasets/
    └── ${dataset_name}/
        └── lmdb/
            ├── train
            │   ├── imgs
            │   └── pairs
            ├── valid
            └── test

# 准备finetune相关配置，详见https://github.com/OFA-Sys/Chinese-CLIP#模型finetune
# 指定机器数 & 卡数
GPUS_PER_NODE=1 # 卡数
WORKER_CNT=1 # 机器数
MASTER_ADDR="localhost"
MASTER_PORT=8514 # 同台机器同时起多个任务，请分别分配不同的端口号
RANK=0

# 刚刚创建过的目录，存放了预训练参数和预处理好的数据集
DATAPATH="../fmhData"
DATASET="Flickr30k-CN"

# 指定LMDB格式的训练集和验证集路径（存放了LMDB格式的图片和图文对数据）
train_data=f"{DATAPATH}/datasets/{DATASET}/lmdb/train"
val_data=f"{DATAPATH}/datasets/{DATASET}/lmdb/valid"
num_workers=4 # 训练集pytorch dataloader的进程数，设置为>0，以减小训练时读取数据的时间开销
valid_num_workers=4 # 验证集pytorch dataloader的进程数，设置为>0，以减小验证时读取数据的时间开销

# 指定刚刚下载好的Chinese-CLIP预训练权重的路径
resume=f"{DATAPATH}/pretrained_weights/clip_cn_vit-h-14.pt"
reset_data_offset="--reset-data-offset" # 从头读取训练数据
reset_optimizer="--reset-optimizer" # 重新初始化AdamW优化器

# 指定输出相关配置
batchsize=64
output_base_dir=f"{DATAPATH}/experiments/"
#name=f"flickr30kcn_finetune_vit-l-14_roberta-base_batchsize{batchsize}_1gpu" # finetune超参、日志、ckpt将保存在../datapath/experiments/
name=f"flickr30kcn_finetune_vit-h-14_roberta-large_batchsize{batchsize}_1gpu" # finetune超参、日志、ckpt将保存在../datapath/experiments/
save_step_frequency=999999 # disable it
save_epoch_frequency=1 # 每轮保存一个finetune ckpt
log_interval=10 # 日志打印间隔步数
report_training_batch_acc="--report-training-batch-acc" # 训练中，报告训练batch的in-batch准确率

# 指定训练超参数
context_length=52 # 序列长度，这里指定为Chinese-CLIP默认的52
warmup=100 # warmup步数
batch_size=batchsize # 训练单卡batch size
valid_batch_size=batchsize # 验证单卡batch size
lr=3e-6 # 学习率，因为这里我们使用的对比学习batch size很小，所以对应的学习率也调低一些
wd=0.001 # weight decay
max_epochs=1 # 训练轮数，也可通过--max-steps指定训练步数
valid_step_interval=1000 # 验证步数间隔
valid_epoch_interval=1 # 验证轮数间隔
#vision_model="ViT-L-14-336" # 指定视觉侧结构为ViT-L/14@336
vision_model="ViT-H-14" # 指定视觉侧结构为ViT-H/14
#text_model="RoBERTa-wwm-ext-base-chinese" # 指定文本侧结构为RoBERTa-base
text_model="RoBERTa-wwm-ext-large-chinese" # 指定文本侧结构为RoBERTa-large
use_augment="--use-augment" # 对图像使用数据增强
grad_checkpointing="--grad-checkpointing" # 激活重计算策略，用更多训练时间换取更小的显存开销

run_command = "export PYTHONPATH=${PYTHONPATH}:`pwd`/cn_clip;" + \
f"""
python3 -m torch.distributed.run --nproc_per_node={GPUS_PER_NODE} --nnodes={WORKER_CNT} --node_rank={RANK} \
      --master_addr={MASTER_ADDR} --master_port={MASTER_PORT} cn_clip/training/main.py \
      --train-data={train_data} \
      --val-data={val_data} \
      --num-workers={num_workers} \
      --valid-num-workers={valid_num_workers} \
      --resume={resume} \
      {reset_data_offset} \
      {reset_optimizer} \
      --logs={output_base_dir} \
      --name={name} \
      --save-step-frequency={save_step_frequency} \
      --save-epoch-frequency={save_epoch_frequency} \
      --log-interval={log_interval} \
      {report_training_batch_acc} \
      --context-length={context_length} \
      --warmup={warmup} \
      --batch-size={batch_size} \
      --valid-batch-size={valid_batch_size} \
      --valid-step-interval={valid_step_interval} \
      --valid-epoch-interval={valid_epoch_interval} \
      --lr={lr} \
      --wd={wd} \
      --max-epochs={max_epochs} \
      --vision-model={vision_model} \
      {use_augment} \
      {grad_checkpointing} \
      --text-model={text_model}
""".lstrip()
print(run_command)

# jupyterlab 执行finetune流程  # torch==2.1.0
!{run_command}

cfgfile = "./configs/swinv2_model_eval_config.pkl"  # eval mode
pthfile = "./models/swinv2/swinv2_base_patch4_window12_192_22k.pth"
with open(cfgfile, "rb") as f:
    config = pickle.load(f)
# conifg.EVAL_MODE == True
if config.EVAL_MODE is True:    # turn-off EVAL mode
    config.defrost()
    config.EVAL_MODE = False
    config.freeze()

model = build_model(config)
checkpoint = torch.load(pthfile, map_location="cpu")
model.load_state_dict(checkpoint["model"], strict=False)

# Assuming num_labels is the number of classes for the multi-label task
NUM_LABELS = 20  # Adjust this based on your specific task: SewerML dataset
# Modify the output layer of the SwinV2 model
model.head = nn.Sequential(
    nn.Linear(model.head.in_features, NUM_LABELS),
    nn.Sigmoid(),  # Apply sigmoid for multi-label classification
)

# Binary Cross Entropy Loss and Optimizer
criterion = nn.BCELoss()  # Use BCE Loss for multi-label classification
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

# Training loop
num_epochs = 4  # Specify the number of epochs
model.to(device)
for epoch in range(num_epochs):
    print(f"Epoch {epoch+1}/{num_epochs}")
    # ....