DiT

赛题说明

问题背景：

当前DiT模型用途广泛，主要用于图像生成任务。基于Transformer架构，DiT模型的计算涉及到了矩阵乘法等线性代数运算，这对于在CPU上训练模型以及推理会带来不小的代价，因此，针对在鲲鹏平台上，对DiT模型推理速度加速优化是一个有挑战性的研究课题。

问题描述：

在执行推理过程时，模型的输入包括： ① Noised Latent（数据大小为[batch, 4, 32, 32]，其中batch为批大小，4为特征数量，32为图片大小）； ② Timestep t（数据大小为[batch,1]）； ③ 数据类别标签Label y（数据大小为[batch,1]），使用DiT-XL/2模型的默认参数，则输出数据大小为： [batch, 8, 32, 32] 。其中， [: , 0:4, : , :] 为Noise， [:, 4: , : , :] 为方差Σ。为更好地评估优化手段，本问题设计如下简化条件：（1）本问题中的batch设置为1，模型在鲲鹏平台单节点环境下进行推理；（2）输入数据中的Noised Latent（图片规格为 32×32）是由imagenet-mini开源数据集（图片规格为 256×256 ）经过VAE编码后得到，输出的Noise（大小为[batch, 4, 32, 32] ）可经过VAE解码得到规格为256×256的图片，见图2，图3。在本问题中不考虑这两个部分的耗时。

任务目标：在/home/data/DiT-main 路径下提供了dit的源码，结合算子融合或SME优化算子等优化手段，给出超越基线的推理优化方案。在给定测试集上需满足如下验证标准（推理输入数据为单张图片，大小为[1, 3, 256, 256] ，输出数据也是图片，大小为[1, 3, 256, 256]（使用DiT模型的输出经过VAE解码得到）

登录集群

# 用户名和探索1000集群一致
ssh thuscc25team00@1.92.123.254

代码路径

# /home/data/DiT-main
cp -r /home/data/DiT-main/ /your/path/DiT-main

环境配置实例

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-aarch64.sh
bash Miniconda3-latest-Linux-aarch64.sh
conda init --all
source ~/.bashrc

pip install torch
pip install diffusers
pip install timm

运行实例

# 进入DiT-main
srun -N 1 python sample.py --ckpt /home/data/DiT-XL-2-256x256.pt 2>&1 | tee output.log

评分规则

基准时间：120 秒
只有正确运行且实际运行时间 T < 120 秒 才能获得性能分，否则得 0 分。

T ≤ 96 s（比基准快 ≥ 20%） → 满分 100 分
96 s < T < 120 s → 线性得分
T ≥ 120 s → 0 分

计算公式

\[ S_{\text{perf}} = \min\left(1,\ \max\left(0,\ \dfrac{120 - T}{24}\right)\right) \]

优化报告与代码

初赛结束时，需要交一份压缩包，命名格式为队伍编号，来帮助我们判断你们这道题的优化的合法性。要求如下：

一份 report.pdf 详细说明优化方法
一份最终版本的源代码
包含提交密码的文本文件

DiT提交链接