先放论文
Unsupervised Pre-Training of Image Features on Non-Curated Data
以及Facebook的项目地址
GitHub: Deeper Cluster

环境

• Ubuntu 18.04
• CUDA 10.0
• Python 3.7
• 单机器单GPU

下载源码及预训练模型

sudo git clone https://github.com/facebookresearch/DeeperCluster.git
cd DeeperCluster
sudo ./download_models.sh

安装APEX

sudo git clone https://github.com/NVIDIA/apex.git
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

安装Faiss

不同CUDA版本安装参见官方文档

conda install faiss-gpu cudatoolkit=10.0 -c pytorch # For CUDA10

代码修改

如果要读取自己的数据集
在import部分添加from src.data.loader import load_data
将main.py中load data的代码改为dataset = load_data(args)

参数设置

新建start.sh，设置参数以适应单机器单GPU情况

mkdir -p ./exp/deepercluster/
export NGPU=1 
python -m torch.distributed.launch --nproc_per_node=$NGPU main.py \
        --dump_path ./exp/deepercluster/ \
        --pretrained ./downloaded_models/deepercluster/ours.pth \
        --data_path /home/data/ \
        --size_dataset 5712 \
        --workers 10 \
        --sobel true \
        --lr 0.1 \
        --wd 0.00001 \
        --nepochs 100 \
        --batch_size 48 \
        --reassignment 3 \
        --dim_pca 4096 \
        --super_classes 1 \
        --rotnet false \
        --k 100 \
        --warm_restart false \
        --use_faiss true \
        --niter 10 \
        --world-size 1 \
        --dist-url 'env://'

可修改的：
--pretrained 如果需要调用不同的预训练模型，修改预训练模型地址
--data_path 你自己的数据集地址，该文件夹下需还有文件夹，再内层才是图片
--size_data 这里选择设置数据集大小为batch_size的倍数，否则会报错，同时数据集中真实图片数量与size一致
--workers 读图时候的线程数？
--lr --nepochs --batch_size 常见参数
--reassignment 每几个epoch重新聚类
--super_classes ？
--rotnet 是否启用rotnet，如果启用则--super_classes必须设置为4的倍数
--k K-means的K

遇到过的报错：

1. IndexError，数组越界问题，将size_data设置为batch_size的倍数即可
2. TypeError: ‘numpy.float64’ object cannot be interpreted as an integer 在GitHub上面一个关闭的issue里面提到了，解决方法为将./src/util.py文件中351行的shape=(data.max() + 1, data.size)) 改为shape=(int(data.max()) + 1, data.size))
3. torch.distributed subprocess.CalledProcessError，升级gcc版本

查看结果

结果保存在了./exp/deepercluster/cluster_assignments0.pkl文件中，记得把源文件目录下的文件sort一下再对应结果文件中的类