【SVG 生成系列论文（九）】如何通过文本生成 svg logo？IconShop 模型推理代码详解

SVG 生成系列论文（一）和 SVG 生成系列论文（二）分别介绍了 StarVector 的大致背景和详细的模型细节。
SVG 生成系列论文（三）和 SVG 生成系列论文（四）则分别介绍实验、数据集和数据增强细节。
SVG 生成系列论文（五）介绍了从光栅图像（如 PNG、JPG 格式）转换为矢量图形（如 SVG、EPS 格式）的关键技术-像素预过滤（pixel prefiltering）， Diffvg 这篇论文也是 SVG 生成与编辑领域中 “基于优化”方法的开创性研究。
SVG 生成系列论文（六）和 SVG 生成系列论文（七）简要介绍了 IconShop 的背景、应用和部分细节。
SVG 生成系列论文（八）则介绍了模型架构和具体的训练技巧。

本文将详细拆解 IconShop（论文原文以及代码🔗）的模型结构和对应开源代码。上篇有提到过模型架构如下所示，本篇则从代码的逻辑进行解释，主要是 /path/to/IconShop/model/decoder.py 中的 sample 以及 forward 两个函数。

模型架构

架构整体分为 4 个部分：SVG 图标嵌入（SVG Icon Embedding），文本嵌入（Text Embedding），输入准备和输出生成。

在这里插入图片描述

sample 函数

在原项目中，调用模型推理是从 sample_pixels = sketch_decoder.sample(n_samples=BS, text=tokenized_text) 中进行的，因此重点解读这部分代码。

输入部分：
self: 指代调用该方法的对象，意味着可以访问类的属性和方法。
n_samples: 要生成的样本数量，这里是 batchsize。
text: 输入的文本数据，用于条件生成。形状为[batch_size, text_len]，其中text_len是文本序列的长度，默认是 50。
pixel_seq: 已有的像素序列，初始化为None，形状为[batch_size, max_len]。
xy_seq: 已有的坐标序列，初始化为None，形状为[batch_size, max_len, 2]。

def sample(self, n_samples, text, pixel_seq=None, xy_seq=None):
    """ sample from distribution (top-k, top-p) """
    pix_samples = []
    xy_samples = []
    # latent_ext_samples = []
    top_k = 0
    top_p = 0.5

初始化:
定义了空列表pix_samples 和 xy_samples 以存储采样的像素和坐标序列。
同时定义了 top_k 和 top_p 作为采样策略的参数，默认使用“Top-P”策略（top_p=0.5），不使用“Top-K”。了解更多，可参见大模型推理常见采样策略(Top-k, Top-p)

    # Mapping from pixel index to xy coordiante
    pixel2xy = {}
    x=np.linspace(0, BBOX-1, BBOX)
    y=np.linspace(0, BBOX-1, BBOX)
    xx,yy=np.meshgrid(x,y)
    xy_grid = (np.array((xx.ravel(), yy.ravel())).T).astype(int)
    for pixel, xy in enumerate(xy_grid):
      pixel2xy[pixel] = xy+COORD_PAD+SVG_END

像素到坐标映射:
创建一个从像素索引到坐标(x, y)的映射。这里假设有一个固定大小的边界框(BBOX=200)，通过网格生成所有可能的坐标组合，并将这些坐标与像素索引关联起来。其中，COORD_PAD = NUM_END_TOKEN + NUM_MASK_AND_EOM（NUM_END_TOKEN = 3，NUM_MASK_AND_EOM = 2），SVG_END = 1
处理文本输入: 确保文本序列的长度不超过预设的最大长度self.text_len
循环采样：

for k in range(text.shape[1] + pixlen, self.total_seq_len): 是开始对文本之后的 SVG 序列进行采样（预测）。
初始化或更新像素pixel_seq和坐标序列xy_seq的长度。
使用模型的forward方法计算当前状态下下一个token的概率分布。
应用“Top-K”和“Top-P”过滤策略（top_k_top_p_filtering）到概率分布上，仅保留最可能的token。
从过滤后的分布中多分类采样得到下一个像素值。
将采样的像素索引转换为实际的坐标。其中 PIX_PAD = NUM_END_TOKEN + NUM_MASK_AND_EOM， SVG_END = 1

序列更新和早停条件：

将生成的像素和坐标序列加入现有序列中。
检查生成是否完成（例如生成到结束符），如果有完成的样本则记录并移除。
如果所有样本都生成完毕，提前停止循环。

返回结果：
返回生成的坐标序列xy_samples。

    # Sample per token
    text = text[:, :self.text_len]
    pixlen = 0 if pixel_seq is None else pixel_seq.shape[1]
    for k in range(text.shape[1] + pixlen, self.total_seq_len):
      if k == text.shape[1]:
        pixel_seq = [None] * n_samples
        xy_seq = [None, None] * n_samples
      
      # pass through model
      with torch.no_grad():
        p_pred = self.forward(pixel_seq, xy_seq, None, text)
        p_logits = p_pred[:, -1, :]

      next_pixels = []
      # Top-p sampling of next pixel
      for logit in p_logits: 
        filtered_logits = top_k_top_p_filtering(logit, top_k=top_k, top_p=top_p)
        next_pixel = torch.multinomial(F.softmax(filtered_logits, dim=-1), 1)
        next_pixel -= self.num_text_token
        next_pixels.append(next_pixel.item())

      # Convert pixel index to xy coordinate
      next_xys = []
      for pixel in next_pixels:
        if pixel >= PIX_PAD+SVG_END:
          xy = pixel2xy[pixel-PIX_PAD-SVG_END]
        else:
          xy = np.array([pixel, pixel]).astype(int)
        next_xys.append(xy)
      next_xys = np.vstack(next_xys)  # [BS, 2]
      next_pixels = np.vstack(next_pixels)  # [BS, 1]
        
      # Add next tokens
      nextp_seq = torch.LongTensor(next_pixels).view(len(next_pixels), 1).cuda()
      nextxy_seq = torch.LongTensor(next_xys).unsqueeze(1).cuda()
      
      if pixel_seq[0] is None:
        pixel_seq = nextp_seq
        xy_seq = nextxy_seq
      else:
        pixel_seq = torch.cat([pixel_seq, nextp_seq], 1)
        xy_seq = torch.cat([xy_seq, nextxy_seq], 1)
      
      # Early stopping
      done_idx = np.where(next_pixels==0)[0]
      if len(done_idx) > 0:
        done_pixs = pixel_seq[done_idx] 
        done_xys = xy_seq[done_idx]
        # done_ext = latent_ext[done_idx]
       
        # for pix, xy, ext in zip(done_pixs, done_xys, done_ext):
        for pix, xy in zip(done_pixs, done_xys):
          pix = pix.detach().cpu().numpy()
          xy = xy.detach().cpu().numpy()
          pix_samples.append(pix)
          xy_samples.append(xy)
          # latent_ext_samples.append(ext.unsqueeze(0))
  
      left_idx = np.where(next_pixels!=0)[0]
      if len(left_idx) == 0:
        break # no more jobs to do
      else:
        pixel_seq = pixel_seq[left_idx]
        xy_seq = xy_seq[left_idx]
        text = text[left_idx]
    
    # return pix_samples, latent_ext_samples
    return xy_samples

forward

这个函数的主要作用是前向传播（或称推理），根据输入的像素序列、坐标序列、掩码和文本，计算模型的输出。具体过程如下：

准备输入：

如果需要计算损失return_loss，则去掉最后一个时间步的数据。
计算上下文序列的长度c_seqlen，包括文本和像素序列的长度。

def forward(self, pix, xy, mask, text, return_loss=False):
    '''
    pix.shape  [batch_size, max_len]
    xy.shape   [batch_size, max_len, 2]
    mask.shape [batch_size, max_len]
    text.shape [batch_size, text_len]
    '''
    pixel_v = pix[:, :-1] if return_loss else pix
    xy_v = xy[:, :-1] if return_loss else xy
    pixel_mask = mask[:, :-1] if return_loss else mask

    c_bs, c_seqlen, device = text.shape[0], text.shape[1], text.device
    if pixel_v[0] is not None:
      c_seqlen += pixel_v.shape[1]

嵌入计算：

对输入的文本进行嵌入计算text_emb。
如果有像素和坐标序列，则分别对它们进行嵌入计算(coord_embed_x, pixel_embed)，并与文本嵌入拼接。


    # Context embedding values
    context_embedding = torch.zeros((1, c_bs, self.embed_dim)).to(device) # [1, bs, dim]

    # tokens.shape [batch_size, text_len, emb_dim]
    tokens = self.text_emb(text)

    # Data input embedding
    if pixel_v[0] is not None:
      # coord_embed.shape [batch_size, max_len-1, emb_dim]
      # pixel_embed.shape [batch_size, max_len-1, emb_dim] 
      coord_embed = self.coord_embed_x(xy_v[...,0]) + self.coord_embed_y(xy_v[...,1]) # [bs, vlen, dim]
      pixel_embed = self.pixel_embed(pixel_v)
      embed_inputs = pixel_embed + coord_embed

      # tokens.shape [batch_size, text_len+max_len-1, emb_dim]
      tokens = torch.cat((tokens, embed_inputs), dim=1)

位置编码和掩码计算：

计算位置编码，并与嵌入后的输入序列拼接。
生成用于Transformer的掩码nopeak_mask，确保模型不会看到未来的时间步。
如果有像素掩码pixel_mask，则将其扩展到与嵌入序列匹配的形状。

    # nopeak_mask.shape [c_seqlen+1, c_seqlen+1]
    nopeak_mask = torch.nn.Transformer.generate_square_subsequent_mask(c_seqlen+1).to(device)  # masked with -inf
    if pixel_mask is not None:
      # pixel_mask.shape [batch_size, text_len+max_len]
      pixel_mask = torch.cat([(torch.zeros([c_bs, context_embedding.shape[0]+self.text_len])==1).to(device), pixel_mask], axis=1)

Transformer解码：

将处理后的输入序列送入Transformer解码器，得到输出序列。

    decoder_out = self.decoder(tgt=decoder_inputs, memory=memory_encode, memory_key_padding_mask=None,tgt_mask=nopeak_mask, tgt_key_padding_mask=pixel_mask)

计算logits和损失：

通过全连接层logit_fc计算输出的logits。
如果需要计算损失，则分别计算文本和像素的损失cross_entropy，并返回总损失。
如果不需要计算损失，则直接返回logits。

    # Logits fc
    logits = self.logit_fc(decoder_out)  # [seqlen, bs, dim] 
    logits = logits.transpose(1,0)  # [bs, textlen+seqlen, total_token] 

    logits_mask = self.logits_mask[:, :c_seqlen+1]
    max_neg_value = -torch.finfo(logits.dtype).max
    logits.masked_fill_(logits_mask, max_neg_value)

    if return_loss:
      logits = rearrange(logits, 'b n c -> b c n')
      text_logits = logits[:, :, :self.text_len]
      pix_logits = logits[:, :, self.text_len:]

      pix_logits = rearrange(pix_logits, 'b c n -> (b n) c')
      pix_mask = ~mask.reshape(-1)
      pix_target = pix.reshape(-1) + self.num_text_token

      text_loss = F.cross_entropy(text_logits, text)
      pix_loss = F.cross_entropy(pix_logits[pix_mask], pix_target[pix_mask], ignore_index=MASK+self.num_text_token)
      loss = (text_loss + self.loss_img_weight * pix_loss) / (self.loss_img_weight + 1)
      return loss, pix_loss, text_loss
    else:
      return logits