图像自回归生成(Auto-regressive image generation)实战学习（五）

相关项目下载链接

本节内容详细解析基于 Transformer 的图像补丁令牌预测与生成。这份代码是适配 PatchAutoEncoder+BSQ 二值量化模块的自回归模型实现，核心是完成图像补丁整数令牌的下一个令牌预测，并支持从空序列开始的逐令牌自回归生成。最终能实现图像补丁令牌的生成，并联动 BSQ+PatchAutoEncoder 还原为图像。具体流程如下：

制作训练/验证数据集

基于前文已经训练好的BSQPatchAutoEncoder模型，首先创建token级别的数据集

python-m homework.tokenize checkpoints/YOUR_BSQPatchAutoEncoder.pth data/tokenized_train.pth data/train/*.jpg python-m homework.tokenize checkpoints/YOUR_BSQPatchAutoEncoder.pth data/tokenized_valid.pth data/valid/*.jpg

这会在data目录下生成两个数据集tokenized_train.pth和tokenized_valid.pth，这两个数据集包含了全部训练集和验证集的token数据。
可以通过如下代码检验生成的文件大小：

du-hs data/tokenized_train.pth

如果前文使用的超参数为patch_size=5和codebook_bits = 10，那么文件的大小约为76Mb。

导入依赖库

importabcimporttorchimporttorch.nn.functionalasF

定义模型加载函数

defload()->torch.nn.Module:frompathlibimportPath model_name="AutoregressiveModel"model_path=Path(__file__).parent/f"{model_name}.pth"print(f"Loading{model_name}from{model_path}")returntorch.load(model_path,weights_only=False)

定义抽象基类

基于abc.ABC定义自回归模型的强制实现接口，约束所有自回归模型必须实现forward（预测）和generate（生成）方法。

classAutoregressive(abc.ABC):@abc.abstractmethoddefforward(self,x:torch.Tensor)->tuple[torch.Tensor,dict[str,torch.Tensor]]:# 输入：整数令牌张量(B, h, w)；输出：令牌概率分布(B, h, w, n_token) + 损失/指标字典passdefgenerate(self,B:int=1,h:int=20,w:int=30,device=None)->torch.Tensor:# 生成：从空序列生成B个(h,w)尺寸的整数令牌，默认设备自动适配pass

定义主模型

主模型对应的代码在autoregressive.py。主模型继承torch.nn.Module和抽象基类Autoregressive，是代码的核心实现部分。共包含模型初始化、因果掩码生成、前向预测、自回归生成四大核心部分。

classAutoregressiveModel(torch.nn.Module,Autoregressive):""" Implement an auto-regressive model. The input is a set of patch tokens (integers), the output is an image of probability. You need to implicitly shift your inputs by one position in the forward pass. Make sure n_tokens matches your BSQ dimension (2**codebook_bits_). Hint: You will need the torch.nn.Embedding function Hint: You can use torch.nn.TransformerEncoderLayer if you'd like Hint: You can complete this homework without using positional embeddings """def__init__(self,d_latent:int=128,n_tokens:int=2**10):super().__init__()self.d_latent=d_latent self.n_tokens=n_tokens self.L_max=1024self.embedding=torch.nn.Embedding(num_embeddings=n_tokens,embedding_dim=d_latent)self.pos_emb=torch.nn.Embedding(num_embeddings=self.L_max,embedding_dim=d_latent)encoder_layer=torch.nn.TransformerEncoderLayer(d_model=d_latent,nhead=8,dim_feedforward=4*d_latent,activation="gelu",batch_first=True,norm_first=True,dropout=0.1)self.transformer=torch.nn.TransformerEncoder(encoder_layer=encoder_layer,num_layers=2,norm=torch.nn.LayerNorm(d_latent))self.fc_out=torch.nn.Linear(d_latent,n_tokens)def_generate_causal_mask(self,L:int,device:torch.device)->torch.Tensor:""" 生成因果掩码：确保序列中第i个位置只能看到前i-1个位置，屏蔽每个位置对后续位置的注意力 :param L: 序列长度 h*w :param device: 设备 :return: 掩码 (L, L)，float型，上三角=-inf，下三角=0 """mask=torch.nn.Transformer.generate_square_subsequent_mask(L,device=device)returnmaskdefforward(self,x:torch.Tensor)->tuple[torch.Tensor,dict[str,torch.Tensor]]:ifx.dim()==4:x=x.squeeze(1)B,h,w=x.shape L=h*w# 展平成序列x_flat=x.reshape(B,L)# 嵌入 + 位置编码token_emb=self.embedding(x_flat)pos_idx=torch.arange(L,device=x.device)pos_emb=self.pos_emb(pos_idx)x_emb=token_emb+pos_emb# 自回归右移（关键）x_emb=F.pad(x_emb,(0,0,1,0))[:,:-1]# 因果掩码mask=self._generate_causal_mask(L,x.device)trans_out=self.transformer(x_emb,mask=mask)# 输出logits=self.fc_out(trans_out)logits_2d=logits.reshape(B,h,w,self.n_tokens)returnlogits_2d,{}@torch.no_grad()defgenerate(self,B:int=1,h:int=30,w:int=20,device=None)->torch.Tensor:# noqapass

模块测评

下面进行模型训练：

python-m homework.train AutoregressiveModel

模型训练过程如图所示

将代码打包为压缩文件

python bundle.py homework20260201

进行评分自测：

python-m grader20260201.zip

最终测试得分如下：