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

相关项目下载链接本节内容详细解析基于 Transformer 的图像补丁令牌预测与生成。这份代码是适配 PatchAutoEncoderBSQ 二值量化模块的自回归模型实现核心是完成图像补丁整数令牌的下一个令牌预测并支持从空序列开始的逐令牌自回归生成。最终能实现图像补丁令牌的生成并联动 BSQPatchAutoEncoder 还原为图像。具体流程如下制作训练/验证数据集基于前文已经训练好的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_size5和codebook_bits 10那么文件的大小约为76Mb。导入依赖库importabcimporttorchimporttorch.nn.functionalasF定义模型加载函数defload()-torch.nn.Module:frompathlibimportPath model_nameAutoregressiveModelmodel_pathPath(__file__).parent/f{model_name}.pthprint(fLoading{model_name}from{model_path})returntorch.load(model_path,weights_onlyFalse)定义抽象基类基于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:int1,h:int20,w:int30,deviceNone)-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 youd like Hint: You can complete this homework without using positional embeddings def__init__(self,d_latent:int128,n_tokens:int2**10):super().__init__()self.d_latentd_latent self.n_tokensn_tokens self.L_max1024self.embeddingtorch.nn.Embedding(num_embeddingsn_tokens,embedding_dimd_latent)self.pos_embtorch.nn.Embedding(num_embeddingsself.L_max,embedding_dimd_latent)encoder_layertorch.nn.TransformerEncoderLayer(d_modeld_latent,nhead8,dim_feedforward4*d_latent,activationgelu,batch_firstTrue,norm_firstTrue,dropout0.1)self.transformertorch.nn.TransformerEncoder(encoder_layerencoder_layer,num_layers2,normtorch.nn.LayerNorm(d_latent))self.fc_outtorch.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 masktorch.nn.Transformer.generate_square_subsequent_mask(L,devicedevice)returnmaskdefforward(self,x:torch.Tensor)-tuple[torch.Tensor,dict[str,torch.Tensor]]:ifx.dim()4:xx.squeeze(1)B,h,wx.shape Lh*w# 展平成序列x_flatx.reshape(B,L)# 嵌入 位置编码token_embself.embedding(x_flat)pos_idxtorch.arange(L,devicex.device)pos_embself.pos_emb(pos_idx)x_embtoken_embpos_emb# 自回归右移关键x_embF.pad(x_emb,(0,0,1,0))[:,:-1]# 因果掩码maskself._generate_causal_mask(L,x.device)trans_outself.transformer(x_emb,maskmask)# 输出logitsself.fc_out(trans_out)logits_2dlogits.reshape(B,h,w,self.n_tokens)returnlogits_2d,{}torch.no_grad()defgenerate(self,B:int1,h:int30,w:int20,deviceNone)-torch.Tensor:# noqapass模块测评下面进行模型训练python-m homework.train AutoregressiveModel模型训练过程如图所示将代码打包为压缩文件python bundle.py homework20260201进行评分自测python-m grader20260201.zip最终测试得分如下