DALLE-pytorch

环境要求

• Python 环境（需能 pip install）
• PyTorch（含 CUDA 版本以便 GPU 训练）
• 若要使用 Docker 训练，需安装 Docker 与 Docker Container Runtime
• GPU（训练 DALL-E 建议使用，docker run 中带 --gpus all）

安装与启动步骤

1. 1安装依赖库

   通过 pip 直接安装 dalle-pytorch 包，这是 README 中给出的唯一安装方式。

   pip install dalle-pytorch

   复制

2. 2克隆仓库（建 Docker）

   若要用官方 Docker 训练，需先获取仓库源码，docker build 需要 docker 目录。

   git clone https://github.com/lucidrains/DALLE-pytorch.git
   cd DALLE-pytorch

   复制

3. 3构建 Docker 镜像

   在仓库根目录执行，镜像名为 dalle，使用 docker 目录下的 Dockerfile 保证 PyTorch/CUDA 版本正确。

   docker build -t dalle docker

   复制

4. 4启动容器交互训练

   挂载当前目录到 /workspace/dalle 并开启 GPU，进入容器后即可运行训练脚本。

   docker run --gpus all -it --mount src="$(pwd)",target=/workspace/dalle,type=bind dalle:latest bash

   复制

5. 5创建 VAE 模型

   在 Python 中导入 DiscreteVAE 并按 README 示例设置图像尺寸、层数、视觉 token 数等参数。

   import torch
   from dalle_pytorch import DiscreteVAE
   
   vae = DiscreteVAE(
       image_size = 256,
       num_layers = 3,
       num_tokens = 8192,
       codebook_dim = 512,
       hidden_dim = 64,
       num_resnet_blocks = 1,
       temperature = 0.9,
       straight_through = False,
   )

   复制

6. 6跑一次 VAE 前向传播

   用随机图像张量验证模型能正常计算损失并反向传播，确认依赖和 GPU 配置无误。

   images = torch.randn(4, 3, 256, 256)
   loss = vae(images, return_loss = True)
   loss.backward()

   复制

关键配置

如何确认成功

执行 vae(images, return_loss=True) 可正常返回 loss 并 loss.backward() 不报错，说明环境与模型配置正确。

常见问题

注意事项

• README 仅给出 VAE 部分的 Python 示例，Transformer 训练与 CLIP 重排序细节需参考官方 wiki
• num_tokens 可按项目规模调小以降低显存占用
• temperature 越低离散化越强，可按需调整
• Docker 运行需确保宿主机已安装 Docker 与 Container Runtime 并支持 GPU