DoRA_ICLR24

DoRA_ICLR24

该仓库包含 ICLR 2024 论文“Is ImageNetworth 1 video? Learning Strong image Encoders from 1 long unlabeled video”的官方实现

[ arXiv ], [ paper ], [ dataset ], [ Project Page ]

概述

动机

我们的目标是通过利用视频帧中的丰富信息来构建强大的表示。标准 SSL 框架（例如 SimCLR、DINo 等）通常假定不同视图之间存在对应关系。无论是通过池化使用密集表示还是全局表示（例如 iBOT），都是如此。虽然在图像中建立对应关系相对简单，但在处理时间变形时变得更具挑战性，需要某种形式的对象跟踪。在具有大视野或自我运动的视频中，获得对应关系变得更加困难。

高层次的想法

我们介绍基于多对象发现和跟踪的DoRA。它利用视觉转换器中不同头部的 [CLS] 标记的注意力来跨时间序列识别并一致地跟踪给定帧内的多个对象。然后对这些应用师生蒸馏损失。重要的是，我们不使用任何现成的对象跟踪器或光流网络。这使我们的管道保持简单，不需要任何额外的数据或培训。它还确保学习到的表示是稳健的。

数据集准备

原创视频

HuggingFace 数据集中提供了 WTour 视频的完整列表，其中包括 YouTube 链接以及有关如何使用必要脚本下载视频的详细说明。

如果链接损坏，您可以直接从徒步旅行数据集下载视频

脸部模糊的视频

为了遵守 GDPR，我们还尝试使用 Deface 模糊视频中出现的所有面孔和车牌

要对 WTour 数据集中的所有视频执行此操作：

python3 -m pip install deface

然后使用 bash 脚本对所有视频运行 Deface：

chmod a+x gdpr_blur_faces.sh  
./gdpr_blur_faces.sh

训练多拉

安装软件包

这些要求可以通过 Anaconda 轻松安装。这里我们创建一个名为dora的conda环境并安装所有必需的库

conda create -n dora python=3.8.2
conda activate dora
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 pytorch-cuda=11.7 -c pytorch -c nvidia
conda install pillow==9.4.0 scipy tensorboardX faiss-gpu==1.6.1 tqdm lmdb scikit-learn pyarrow==2.0.0 DALL-E munkres six einops

预训练 DoRA

如果您使用具有 32GB RAM 的 V100 GPU，我们建议使用以下脚本来预训练 DoRA：

python -m torch.distributed.launch --nproc_per_node=8 main.py --arch vit_small --data_path /path-to-video/venice.mp4 
  --output_dir /path-to-checkpoint/weights_WTours/DoRA/venice_8frames/ --optimizer adamw 
  --use_bn_in_head False --out_dim 65536 --batch_size_per_gpu 6 --local_crops_number 6 --epochs 100  
  --num_workers 10 --lr 0.0005 --min_lr 0.00001  --norm_last_layer False  
  --warmup_teacher_temp_epochs 30 --weight_decay 0.04 --weight_decay_end 0.4 
  --frame_per_clip 8 --step_between_clips 60

对于具有 80GB RAM 的 A100 GPU，我们建议使用以下预训练脚本（唯一的区别在于batch_size_per_gpu）：

python -m torch.distributed.launch --nproc_per_node=8 main.py --arch vit_small --data_path /path-to-video/venice.mp4 
  --output_dir /path-to-checkpoint/weights_WTours/DoRA/venice_8frames/ --optimizer adamw 
  --use_bn_in_head False --out_dim 65536 --batch_size_per_gpu 16 --local_crops_number 6 --epochs 100  
  --num_workers 10 --lr 0.0005 --min_lr 0.00001  --norm_last_layer False  
  --warmup_teacher_temp_epochs 30 --weight_decay 0.04 --weight_decay_end 0.4 
  --frame_per_clip 8 --step_between_clips 60

预训练模型

您可以下载用于下游任务的预训练主干网的权重。所有主干网均为ViT-S/16型号。

纪元	模型	预训练数据集	下载
100	恐龙	WT-威尼斯	骨干
100	多拉	WT-威尼斯	骨干
300	多拉	WT-威尼斯	骨干
100	多拉	WT-全部	骨干

评估

线性探测和 KNN

用于 ImageNet-1K 上的线性探测

python -m torch.distributed.launch --nproc_per_node=4 eval_linear.py 
--batch_size_per_gpu 512 --n_last_blocks 4 --avgpool_patchtokens 0 --arch vit_small --lr 0.01  
--pretrained_weights /path-to-checkpoint/venice/checkpoint.pth --data_path /dataset/imagenet/ 
--output_dir /path-to-output-dir/venice/LP/

用于 ImageNet-1K 上的 KNN 评估

python3 -m torch.distributed.launch --nproc_per_node=4 eval_knn.py 
--arch vit_small --checkpoint_key teacher --data_path /dataset/imagenet/ 
--pretrained_weights /path-to-checkpoint/venice/checkpoint.pth

语义分割 (ADE20K) 和对象检测 (MS-COCO)

请遵循 iBOT 的评估脚本

视频对象分割 (DAVIS) 和对象跟踪 (GOT-10K)

（待办事项）

奖状

如果您发现此存储库有用，请考虑给予星号和引用：

@inproceedings{venkataramanan2024imagenet,  
  title={Is ImageNet worth 1 video? Learning strong image encoders from 1 long unlabelled video},  
  author={Venkataramanan, Shashanka and Rizve, Mamshad Nayeem and Carreira, Jo{~a}o and Asano, Yuki M and Avrithis, Yannis},  
  booktitle={International Conference on Learning Representations},  
  year={2024}  
}