PyTorch实现MAR+DiffLoss项目教程

PyTorch实现MAR+DiffLoss项目教程

mar PyTorch implementation of MAR+DiffLoss https://arxiv.org/abs/2406.11838 项目地址: https://gitcode.com/gh_mirrors/mar6/mar

1. 项目介绍

本项目是基于PyTorch框架的MAR+DiffLoss的官方实现，它是一种无需向量量化的自回归图像生成方法。MAR（Memory Augmented Autoregressive）模型通过引入记忆机制，改进了传统自回归模型在图像生成中的表现。DiffLoss则是一种损失函数，用于优化生成图像的质量。本项目包含了预训练的模型、训练和评估脚本，以及一个交互式可视化演示。

2. 项目快速启动

环境准备

首先，您需要克隆项目仓库并创建一个合适的conda环境：

git clone https://github.com/LTH14/mar.git
cd mar
conda env create -f environment.yaml
conda activate mar

模型下载

接下来，下载预训练的VAE和MAR模型：

python util/download.py

运行演示

您可以通过以下命令运行交互式可视化演示：

python demo/gradio_app.py

训练模型

以下是启动默认设置（MAR-L，DiffLoss MLP 3个块，宽度为1024通道，400个周期）的命令：

torchrun --nproc_per_node=8 --nnodes=4 --node_rank=${NODE_RANK} --master_addr=${MASTER_ADDR} --master_port=${MASTER_PORT} \
main_mar.py \
--img_size 256 --vae_path pretrained_models/vae/kl16.ckpt --vae_embed_dim 16 --vae_stride 16 --patch_size 1 \
--model mar_large --diffloss_d 3 --diffloss_w 1024 \
--epochs 400 --warmup_epochs 100 --batch_size 64 --blr 1.0e-4 --diffusion_batch_mul 4 \
--output_dir ${OUTPUT_DIR} --resume ${OUTPUT_DIR} \
--data_path ${IMAGENET_PATH}

请根据您的实际情况替换${NODE_RANK}, ${MASTER_ADDR}, ${MASTER_PORT}, ${OUTPUT_DIR}, 和${IMAGENET_PATH}等环境变量。

3. 应用案例和最佳实践

本项目提供了一些预训练的模型，您可以直接用于图像生成。以下是一些应用案例：

• 使用预训练的MAR模型生成指定类别的图像。
• 通过调整--cfg和--temperature参数，使用分类器自由引导生成更高质量的图像。
• 在不同的数据集上微调预训练的模型，以适应特定的应用场景。

4. 典型生态项目

本项目是基于以下开源项目构建的：

• MAE (Masked Autoencoder)：一种自编码器模型，用于图像的表示学习。
• MAGE (Memory Augmented Generative Encoder)：一种引入记忆机制的生成模型。
• DiT (Denoising Diffusion Trees)：一种基于去噪扩散的图像生成模型。

通过结合这些项目的优势，本项目实现了更高效的图像生成算法。

mar PyTorch implementation of MAR+DiffLoss https://arxiv.org/abs/2406.11838 项目地址: https://gitcode.com/gh_mirrors/mar6/mar