这两天我一直在尝试复现 MAE，与我的数据进行结合，来做出一点新的东西。
Offical MAE
我首先从 meta 的官方仓库开始。我成功跑通了 demo，但是发现仍然有些问题。submitit 是 meta 开源的集群任务提交工具。这个代码使用的应该是 meta 内部使用的训练集群，基于 slurm。我无法访问这个集群。而且部署slurm 集群可能很麻烦。因此作罢。
Tips
使用开源工具是对的。内部的私有工具无法迁移到下一个团队，重新学习的成本比较高。开源工具一般没有这个限制，可以打通切换团队后的协作。
Dockerfile
1
2
3
4
5
6
7
8
9
| FROM pytorch/pytorch:1.11.0-cuda11.3-cudnn8-devel
RUN pip install jupyterlab "numpy<1.20" "timm==0.4.5" matplotlib -i https://mirrors.sustech.edu.cn/pypi/simple
RUN rm /etc/apt/sources.list.d/cuda.list && \
rm /etc/apt/sources.list.d/nvidia-ml.list && \
apt-get update && apt-get install -y wget
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token='abcd'"]
|
使用上面这个 Dockerfile 可以跑通环境。
MAE-pytorch
之后我想使用 pytorch lightning 来跑一下 MAE。现在训练的显卡够用，但是 pytorch 直接写 DDP 比较麻烦。于是我参考了 mae-pytorch。
但是训练之后发现，似乎对 lightning 的适配不是很好。我调整了 num_gpus 参数，训练仍然不顺。
https://gist.github.com/Svtter/f4416107d99d69d3a3c48a6eae918411
1
2
3
4
5
6
7
8
9
| File /workspace/mae.py:145, in MaskedAutoencoderLIT.generate_grid_of_samples(self, true_images, pred_patch, mask)
144 pred_images = res[sample_indices][:, [3, 2, 1]]
--> 145 true_image_grid = true_images[sample_indices][:, [3, 2, 1]]
147 grid_images = torch.cat([true_image_grid, pred_images], dim=0)
RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.
|
ImageNet 的问题
而且 ImageNet 我本地没有相关的数据，重新下载也是比较麻烦。（虽然可以用 pretrain weights）。
我下载了 tiny-imagenet，但是没能起到作用。
我尝试了一下，能跑通 MNIST + resize，也就是 Tutorial.ipynb 的内容。因为并行训练的重新调优比较麻烦，我暂时放下，再去看看其他的库。
ImageNet 体积比较大，大家一般会使用小型 imagenet 来做一些测试。例如 yolo-v5 issue 中所表述的。
imagenet-100
imagenet 往往太大了，我们可用一些小型的数据集去验证模型；但是 shape 和文件结构需要与 imagenet 一致。
1
2
3
4
5
6
| import kagglehub
# Download latest version
path = kagglehub.dataset_download("ambityga/imagenet100")
print("Path to dataset files:", path)
|
Dockerfile
1
2
3
4
5
6
7
8
9
| FROM pytorch/pytorch:2.4.1-cuda11.8-cudnn9-devel
WORKDIR /workspace
COPY requirements.txt .
RUN pip install -r requirements.txt -i https://mirrors.sustech.edu.cn/pypi/web/simple && \
pip install jupyterlab matplotlib
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token='abcd'"]
|
solo learn
solo-learn 是发表在 IJML 的一个工作。主要实现了 self-regression 的相关方法。他们使用了一个我之前没有使用过的库NVIDIA/DALI
我使用了 MAE_pytorch 的 Dockerfile，做了一些调整。
Start
做分类任务，参考文档 offline-linear-eval.
首先是文档有点问题，from solo.utils ... 应该改为 from solo.data...。
这个库使用了 hydra 以及 omegaconf。
hydra 可以比较有效的去处理配置重叠的情况。我之前写类似的东西，都是一个一个类，然后不停的嵌套嵌套。这样看，hydra 确实方便很多。
omegaconf 是一个使用 hydra 配置文件的工具；由于 hydra 配置返回的是 DictConfig，因此读取配置的时候需要使用 dict 数据结构进行访问。
不过我觉得实际上用处不是很大。除非项目很庞大，否则增加了上手的成本。
solo-learn 基本上已经不太维护了。我尝试下载了 imagenet-100，正在测试中。
尝试执行了
1
2
3
4
5
6
7
8
| #!/bin/bash
python3 main_pretrain.py \
--config-path scripts/pretrain/imagenet-100/ \
--config-name barlow.yaml
# add new arguments (e.g. those not defined in the yaml files)
# by doing ++new_argument=VALUE
# pytorch lightning's arguments can be added here as well.
|
遇到 cuda 内存不够用的情况。放弃了。
Dockerfile
1
2
3
4
5
6
7
8
9
10
| FROM pytorch/pytorch:2.4.1-cuda11.8-cudnn9-devel
WORKDIR /workspace
COPY requirements.txt .
RUN pip install -r requirements.txt -i https://mirrors.sustech.edu.cn/pypi/web/simple && \
pip install jupyterlab matplotlib
# RUN pip3 install .[dali,umap,h5] --extra-index-url https://developer.download.nvidia.com/compute/redist
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8888", "--no-browser", "--allow-root", "--NotebookApp.token='abcd'"]
|
Tips
直接使用 checkpoints 可以节省大量的训练时间。手里的卡更多的应该去用来 fine tune，或者训练自己的数据集。从头开始训练（我之前的做法）往往是不经济且速度慢的。
很出名的一个库。很多朋友都使用这个库来尝试了构建了自己的大型模型。
我想 MAE 应该是组成部分之一，我想借着这个机会使用一下看。
我尝试了一下 open sources models by huggingface。
transformers 使用起来比较简单。zero-shot image classification 入门速度很快。
使用 vit_mae，可以通过这个官方文档。
使用预训练权重的 MAE
这是 huggingface 官方提供的源代码，我直接贴在这里了。
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
| from transformers import ViTFeatureExtractor, ViTForImageClassification
import torch
from datasets import load_dataset
dataset = load_dataset("huggingface/cats-image")
image = dataset["test"]["image"][0]
feature_extractor = ViTFeatureExtractor.from_pretrained("google/vit-base-patch16-224")
model = ViTForImageClassification.from_pretrained("google/vit-base-patch16-224")
inputs = feature_extractor(image, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
# model predicts one of the 1000 ImageNet classes
predicted_label = logits.argmax(-1).item()
print(model.config.id2label[predicted_label])
|
