CogView：40亿参数的文本到图像生成模型

THUDM 团队开源的 CogView，是一个专注于文本到图像生成的预训练模型，目前已获得 1,797 个 Star：

CogView 基于 40 亿参数的 Transformer 架构，支持根据文本描述生成对应图像，对中文输入有较好的支持。该项目由清华大学知识工程实验室（THUDM）开发，相关论文已被 NeurIPS 2021 接收。

核心能力

CogView 提供三种主要功能：

文本生成图像。用户输入描述文本，模型生成对应图像。支持通过交互式 CLI 或批量文件输入两种方式调用。

图像超分辨率。对已生成的图像进行分辨率提升，基于专用的超分辨率模型处理。

图像到文本。输入图像，模型输出对应的文本描述。该功能目前处于实验阶段，效果尚可但尚未专门优化。

技术细节

CogView 的训练和推理涉及几个关键组件。图像分词器（Image Tokenizer）将图像编码为离散token，预训练模型在此基础上学习文本与图像之间的映射关系。训练过程中采用的 PB-relax 和 Sandwich-LN 技术，可帮助稳定大规模 Transformer 的训练，减少 NaN 损失的出现。

环境要求

运行 CogView 的完整功能需要具备以下硬件和软件条件：

Linux 服务器搭配 Nvidia V100 或 A100 GPU 为推荐配置。显存不足时，可通过调小 --max-inference-batch-size 参数来适配较低规格的显卡。

软件依赖方面，需要 PyTorch 1.7.0 或以上版本，以及 NVIDIA apex 库。项目提供了 requirements.txt 和 Docker 镜像两种环境准备方式。

快速上手

文本生成图像的使用流程如下：

将待生成的文本描述写入 input.txt，每行一条。执行 ./scripts/text2image.sh --debug，结果保存在 samples_text2image/ 目录下。--debug 参数会将同一批生成的样本拼接为一张图片，按输入文本和日期命名。

如需生成独立的分图，使用 --with-id 参数。此时输入格式为 id{tab}文本描述，结果按 id 分目录存储。

超分辨率功能通过 ./scripts/super_resolution.sh 调用，输入格式为 文本{tab}图片路径。图像到文本功能通过 ./scripts/image2text.sh 调用，输入为每行一个图片路径。

模型获取

项目提供三种预训练模型：基础文本到图像模型、图像描述模型、超分辨率模型。用户可从清华云盘或智源研究院的悟道平台下载。同时项目也提供了小型数据集，供训练流程的测试和验证使用。

训练支持

CogView 支持单节点和多节点分布式训练。单节点场景下直接运行 ./scripts/pretrain_single_node.sh 即可。多节点训练需要各节点通过 infiniband 互联，配合 Docker 环境和 pdsh 工具完成部署。

除了本地部署，项目还提供了 Google Colab 运行方案，以及在线 Demo 页面，方便快速体验生成效果。

署，项目还提供了 Google Colab 运行方案，以及在线 Demo 页面，方便快速体验生成效果。