图像作为点集 - ICLR'23 [口头报告，前5%]

作者：Xu Ma*、Yuqian Zhou*、 Huan Wang、Can Qin、Bin Sun、 Chang Liu、Yun Fu。

---

聚类过程（从头到最后一轮）：

---

待办事项（3月9日）：

• 2024年2月24日：用更好的结果重新训练所有模型，并更新了谷歌云盘链接（smilelab网页已失效）。
• 发布不含区域划分的CoC代码和检查点（使用更新后的代码重新训练，效果更好）。
• 发布可视化脚本。

图像分类

1. 需求

torch>=1.7.0；torchvision>=0.8.0；pyyaml；timm==0.6.13；einops；apex-amp（如果想使用fp16）；

数据准备：ImageNet需按以下文件夹结构组织，可使用此脚本解压ImageNet。

│imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......

2. 预训练的Context Cluster模型

我们更新了新的检查点和日志（修复了bug，改进了全连接层实现）。欢迎下载。

模型	参数量	图像分辨率	Top1准确率	吞吐量	下载（更新后的链接，准确率更高）
ContextCluster-tiny	5.3M	224	73.2	518.4	[检查点与日志]
ContextCluster-tiny_plain（无区域划分）	5.3M	224	73.8	-	[检查点与日志]
ContextCluster-small	14.0M	224	78.1	513.0	[检查点与日志]
ContextCluster-medium	27.9M	224	81.2	325.2	[检查点与日志]

3. 验证

要评估我们的Context Cluster模型，请运行：

MODEL=coc_tiny #{tiny, tiny2 small, medium}
python3 validate.py /path/to/imagenet  --model $MODEL -b 128 --checkpoint {/path/to/checkpoint} 

4. 训练

我们展示了如何在8张GPU上训练Context Cluster。学习率与批量大小的关系为lr=bs/1024*1e-3。 为方便起见，假设批量大小为1024，则学习率为1e-3（对于1024的批量大小，有时将学习率设为2e-3会取得更好的效果）。

MODEL=coc_tiny # coc variants
DROP_PATH=0.1 # drop path rates
python3 -m torch.distributed.launch --nproc_per_node=8 train.py --data_dir /dev/shm/imagenet --model $MODEL -b 128 --lr 1e-3 --drop-path $DROP_PATH --amp

5. 聚类可视化

我们提供了一个脚本，用于可视化给定阶段、块和头的CoC聚类结果。 不同层/头会呈现出不同的聚类模式。

# 使用示例（生成的图像将保存到images/cluster_vis/{model}）：
python cluster_visualize.py --image {path_to_image} --model {model} --checkpoint {path_to_checkpoint} --stage {stage} --block {block} --head {head}
 

请参阅pointcloud文件夹，了解ScanObjectNN上的点云分类任务。

请参阅detection文件夹，了解COCO上的目标检测和实例分割任务。

请参阅segmentation文件夹，了解ADE20K上的语义分割任务。

BibTeX

@inproceedings{ma2023image,
    title={Image as Set of Points},
    author={Xu Ma and Yuqian Zhou and Huan Wang and Can Qin and Bin Sun and Chang Liu and Yun Fu},
    booktitle={The Eleventh International Conference on Learning Representations},
    year={2023},
    url={https://openreview.net/forum?id=awnvqZja69}
}

致谢

我们的实现主要基于以下代码库。我们衷心感谢这些作者的杰出工作。

pointMLP、poolformer、pytorch-image-models、mmdetection、mmsegmentation。

许可证

Context Cluster的大部分内容采用Apache License 2.0授权。

Context-Cluster 快速上手指南

Context-Cluster (CoC) 是一种将图像视为点集进行聚类处理的创新视觉模型，发表于 ICLR 2023 (Oral)。本指南帮助开发者快速完成环境配置、模型验证及基础训练。

1. 环境准备

系统要求

• Python: 建议 3.8+
• GPU: 支持 CUDA 的 NVIDIA 显卡（训练建议多卡，如 8x GPU）
• 数据集: ImageNet (需按特定文件夹结构整理)

依赖安装

请确保已安装 PyTorch 和 torchvision，然后安装其他必要依赖。国内用户推荐使用清华或阿里镜像源加速安装。

# 基础依赖
pip install pyyaml timm==0.6.13 einops -i https://pypi.tuna.tsinghua.edu.cn/simple

# 可选：如需使用混合精度训练 (fp16)，请安装 apex
# 注意：apex 需要从源码编译，且需匹配对应的 CUDA 版本
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

核心依赖版本要求：

• torch >= 1.7.0
• torchvision >= 0.8.0
• timm == 0.6.13

数据准备

ImageNet 数据集需整理为以下目录结构：

│imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   └── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   └── ......

2. 安装步骤

本项目无需复杂的 setup.py 安装过程，直接克隆仓库即可使用。

# 克隆项目代码
git clone https://github.com/ma-xu/Context-Cluster.git
cd Context-Cluster

# 下载预训练模型
# 访问官方 Google Drive 链接下载所需模型的 checkpoint 和 logs
# Tiny 模型示例链接：https://drive.google.com/drive/folders/155Lfofxr7vpdcUxL8ZvzlrNX4cUlhn3J?usp=sharing
# 将下载的 .pth 文件放置在项目根目录或指定路径

提示：由于官方仓库托管在 GitHub 且模型在 Google Drive，国内用户若下载缓慢，可自行寻找国内镜像站或使用代理工具。

3. 基本使用

模型验证 (Validation)

使用预训练权重在 ImageNet 验证集上评估模型性能。以下以 coc_tiny 模型为例：

MODEL=coc_tiny
# 请将 /path/to/imagenet 替换为实际数据路径，/path/to/checkpoint 替换为下载的权重文件路径
python3 validate.py /path/to/imagenet --model $MODEL -b 128 --checkpoint /path/to/checkpoint

模型训练 (Training)

在 8 张 GPU 上启动训练。学习率与批次大小的关系为 lr = bs / 1024 * 1e-3。以下示例假设总批次大小为 1024 (每张卡 128)，学习率设为 1e-3。

MODEL=coc_tiny
DROP_PATH=0.1

# 使用 torch.distributed.launch 启动多卡训练
# --amp 表示启用自动混合精度训练以节省显存并加速
python3 -m torch.distributed.launch --nproc_per_node=8 train.py \
    --data_dir /dev/shm/imagenet \
    --model $MODEL \
    -b 128 \
    --lr 1e-3 \
    --drop-path $DROP_PATH \
    --amp

聚类可视化 (Visualization)

查看模型在特定阶段（stage）、模块（block）和头（head）的聚类效果。生成的图片将保存在 images/cluster_vis/{model} 目录下。

python cluster_visualize.py \
    --image {path_to_image} \
    --model {model} \
    --checkpoint {path_to_checkpoint} \
    --stage {stage} \
    --block {block} \
    --head {head}

其他任务：如需进行点云分类、目标检测或语义分割任务，请分别参考项目中的 pointcloud/, detection/, segmentation/ 文件夹。