feizaipp Blog

transformer 在 CV 中的应用(三) ViT 分类网络

Mon, 10 May 2021 00:00:00 +0000

0. 参考资料

transformer 在 CV 中的应用(一) Transformer 介绍

1. 网络结构

ViT 网络是谷歌在 2020 年提出的基于纯 Transformer 实现的分类网络，它完全抛弃了 CNN 网络。 ViT 网络中的 Transformer 与传统意义上的 Transformer 存在明显的不同，传统的 Transformer 是用于 NLP 中的机器翻译任务中，它的结构由 Encoder 和 Decoder 两部分组成，因为要将输入的序列通过 Encoder 网络进行编码，然后将编码后的序列通过 Decoder 解码，最终得到目标语言。但是在视觉任务中，我们只需要对图像进行特征提取，然后将特征通过全链接层输出目标类别。

我们知道，在 NLP 模型中 Transformer 的输入是一个序列，那么对于图像数据， ViT 是将一张图像分成一个个小的 patch ，然后对这些 patch 进行编码。

ViT 网络的特征提取是使用了一个叫做 class token 的结构，该结构是可学习的。它与图像编码后的张量在 dim=1 处进行 cat 操作得到一个新的张量，新张量的第一个元素就是 class token ，这个新的张量与位置编码进行相加后输入到 ViT 网络中，最后学习到的 class token 即为图像的特征图，将这个特征通过全链接层得到类别的输出。

ViT 网络的整体结构如下图所示。

2. 代码解析

ViT 代码整体结构也比较简单，主要是注意张量在网络中传播的过程中维度的变化。下面我们从创建 ViT 网络开始。

ViT 的代码实现中，没有使用像 view 、 reshape 等的函数进行维度的变换，而是使用 einops 库，其实也很简单，看下面代码注释就很好明白了。

from einops import rearrange, repeat
from einops.layers.torch import Rearrange

首先要查创建一个 ViT 类
dim: patch embedding 的维度
depth: Transformer 结构中 Encoder 的个数
heads: MSA 的 head 个数
mlp_dim: Transformer 结构中 FFN 的输出维度
img: 随机初始化一个维度为 [4, 3, 256, 256] ， 4 表示 batch size
mask: 可选的 mask ，用于委托哪个 patch 去使用 attend

preds: 调用 ViT 的前向传播函数，输出预测的类别，维度为 [4, 1000]

if __name__ == "__main__":
  v = ViT(
      image_size = 256,
      patch_size = 32,
      num_classes = 1000,
      dim = 1024,
      depth = 6,
      heads = 16,
      mlp_dim = 2048,
      dropout = 0.1,
      emb_dropout = 0.1
  )

  img = torch.randn(4, 3, 256, 256)
  mask = torch.ones(1, 8, 8).bool() # optional mask, designating which patch to attend to

  # [batch, 1000]
  preds = v(img, mask = mask)

下面看一下 ViT 类的实现，注释中标注了张量维度的变化。

图像的大小是 [256, 256] ，每个 patch 的大小是 [32, 32] ，所以图像被分割成了 64 个 patch

class ViT(nn.Module):
  def __init__(self, *, image_size, patch_size, num_classes, dim, depth, heads, mlp_dim, pool = 'cls', channels = 3, dim_head = 64, dropout = 0., emb_dropout = 0.):
      super().__init__()
      assert image_size % patch_size == 0, 'Image dimensions must be divisible by the patch size.'
      num_patches = (image_size // patch_size) ** 2
      patch_dim = channels * patch_size ** 2
      assert pool in {'cls', 'mean'}, 'pool type must be either cls (cls token) or mean (mean pooling)'

      # img:[4,3,256,256]
      # patch_size=32 ，每个 patch 的大小是 (32, 32)
      # h=8
      # w=8
      # Rearrange: [4, 64, (32*32*3)]=[4, 64, 3072]
      # patch_dim: 3072
      # dim: 1024
      # 经过全链接层后，维度由 [4, 64, 3072] -> [4, 64, 1024]
      self.to_patch_embedding = nn.Sequential(
          Rearrange('b c (h p1) (w p2) -> b (h w) (p1 p2 c)', p1 = patch_size, p2 = patch_size),
          nn.Linear(patch_dim, dim),
      )

      # 位置编码
      # num_patches: (256//32) ** 2=64
      # (1, 65, 1024)
      self.pos_embedding = nn.Parameter(torch.randn(1, num_patches + 1, dim))
      # (1,1,1024)
      self.cls_token = nn.Parameter(torch.randn(1, 1, dim))
      # emb_dropout=0.1
      self.dropout = nn.Dropout(emb_dropout)

      # dim=1024
      # depth=6
      # heads=16
      # dim_head=64
      # mlp_dim=2048
      self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, dropout)

      self.pool = pool
      # 这个函数建立一个输入模块，什么都不做，通常用在神经网络的输入层。这个可以用在残差学习中。
      self.to_latent = nn.Identity()

      # dim=1024
      self.mlp_head = nn.Sequential(
          nn.LayerNorm(dim),
          nn.Linear(dim, num_classes)
      )

  def forward(self, img, mask = None):
      # img:[b,3,256,256]
      # x:[b,64,1024]
      x = self.to_patch_embedding(img)
      b, n, _ = x.shape

      # self.cls_token: [1, 1, dim]
      # cls_tokens: [b, 1, dim]
      cls_tokens = repeat(self.cls_token, '() n d -> b n d', b = b)
      # x: [b, 65, 1024]
      x = torch.cat((cls_tokens, x), dim=1)
      # x 加上位置编码
      x += self.pos_embedding[:, :(n + 1)]
      x = self.dropout(x)

      # x: 输入输出维度都是 [b, 65, 1024]
      x = self.transformer(x, mask)

      # [b, 1024]
      x = x.mean(dim = 1) if self.pool == 'mean' else x[:, 0]

      x = self.to_latent(x)
      return self.mlp_head(x)

下面我们看下 Transformer 的代码， ViT 中的 Transformer 结构中有 Attention 、FeedForward 两个结构。注意此处没有 Decoder 结构，只有 Encoder 。

# dim=1024
# depth=6
# heads=16
# dim_head=64
# mlp_dim=2048
class Transformer(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                Residual(PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout))),
                Residual(PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout)))
            ]))
    def forward(self, x, mask = None):
        for attn, ff in self.layers:
            x = attn(x, mask = mask)
            x = ff(x)
        return x

先看 Attention 。

class Attention(nn.Module):
    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0.):
        super().__init__()
        inner_dim = dim_head *  heads
        # project_out: True
        project_out = not (heads == 1 and dim_head == dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        # 计算 q k v 矩阵
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        ) if project_out else nn.Identity()

    def forward(self, x, mask = None):
        # x: [b, 65, 1024]
        # h:16
        b, n, _, h = *x.shape, self.heads
        # chunk: 将 tensor 进行分割成 3 份，如果指定轴的元素被 3 除不尽，那么最后一块的元素个数变少
        # qkv: [b, 65, inner_dim * 3]
        qkv = self.to_qkv(x).chunk(3, dim = -1)
        # q k v: [b, 65, inner_dim]=[b, 65, 16*64]->[b, 16, 65, 64]
        # map: 取出 qkv 中的每个值，然后执行 lambda 表达式
        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = h), qkv)

        # dots: 返回的维度为 [b, 16, 65, 65] ，相当于 q*k.T
        dots = einsum('b h i d, b h j d -> b h i j', q, k) * self.scale
        # 获取最小值，也就相当于 0
        mask_value = -torch.finfo(dots.dtype).max

        # mask: [1, 8, 8]
        if mask is not None:
            # pad: 矩阵填充函数
            # input: 需要扩充的 tensor ，可以是图像数据，抑或是特征矩阵数据
            # pad: 扩充维度，用于预先定义出某维度上的扩充参数
            # mode: 扩充方法，’constant‘, ‘reflect’ or ‘replicate’三种模式，分别表示常量，反射，复制
            # value: 扩充时指定补充值，但是 value 只在 mode='constant’ 有效，即使用 value 填充在扩充出的新维度位置，而在’reflect’和’replicate’模式下，value不可赋值
            # pad 定义：
            # 如果参数pad只定义两个参数，表示只对输入矩阵的最后一个维度进行扩充
            # 如果参数pad只定义四个参数，前两个参数对最后一个维度有效，后两个参数对倒数第二维有效。
            # 如果参数pad定义六个参数，前4个参数完成了在高和宽维度上的扩张，后两个参数则完成了对通道维度上的扩充。
            # p1d = (左边填充数, 右边填充数)
            # p2d = (左边填充数， 右边填充数， 上边填充数， 下边填充数)
            # p3d = (左边填充数， 右边填充数， 上边填充数， 下边填充数， 前边填充数，后边填充数)
            mask = F.pad(mask.flatten(1), (1, 0), value = True)
            assert mask.shape[-1] == dots.shape[-1], 'mask has incorrect dimensions'
            a = rearrange(mask, 'b i -> b () i ()')
            b = rearrange(mask, 'b j -> b () () j')
            # [b, 1, 65, 1] * [b, 1, 1, 65]=[b, 1, 65, 65]
            mask = rearrange(mask, 'b i -> b () i ()') * rearrange(mask, 'b j -> b () () j')
            dots.masked_fill_(~mask, mask_value)
            del mask

        attn = dots.softmax(dim=-1)

        # out = attn.v
        # [b, 16, 65, 65]*[b, 16, 65, 64]->[b, 16, 65, 64]
        out = einsum('b h i j, b h j d -> b h i d', attn, v)
        out = rearrange(out, 'b h n d -> b n (h d)')
        out =  self.to_out(out)
        return out

再看 FeedForward 。

FFN: 两个全链接层，第一个使用激活函数，第二个不使用。

class FeedForward(nn.Module):
  def __init__(self, dim, hidden_dim, dropout = 0.):
      super().__init__()
      self.net = nn.Sequential(
          nn.Linear(dim, hidden_dim),
          nn.GELU(),
          nn.Dropout(dropout),
          nn.Linear(hidden_dim, dim),
          nn.Dropout(dropout)
      )
  def forward(self, x):
      return self.net(x)

残差网络的实现，先执行 fn 操作，然后与输入进行相加。

class Residual(nn.Module):
    def __init__(self, fn):
        super().__init__()
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(x, **kwargs) + x

归一化的实现，先对输入进行归一化，然后执行 fn 操作

class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn
    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)

transformer 在 CV 中的应用(二) DETR 目标检测网络

Tue, 27 Apr 2021 00:00:00 +0000

0. 参考资料

1. 概述

DETR 目标监测网络是 Facebook 提出的目标监测网络，它是 transformer 在目标检测网络中的首次尝试。相比之前使用 anchor 的目标检测网络， DETR 是一种 anchor free 的目标检测网络，网络最终输出为无序预测集合 (set prediction) 。作者认为像 Faster-RCNN 网络这种设置一大堆的 anchor ，然后基于 anchor 进行分类和回归是属于代理做法，而目标检测任务应该是输出无序集合。那么 DETR 网络结构到底是怎样？它又是如何训练的呢？本文主要介绍 DETR 网络结构以及网络训练流程。

2. 网络结构

DETR 的网络结构由两部分组成， backbone 和 transformer 。

if __name__ == '__main__':
    main(args)

def main(args):
    model, criterion, postprocessors = build_model(args)

def build_model(args):
    return build(args)

def build(args):
    backbone = build_backbone(args)
    transformer = build_transformer(args)
    model = DETR(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        aux_loss=args.aux_loss,
    )

在介绍网络结构之前先介绍下 DETR 网络输入数据的打包格式 NestedTensor 。

该函数输入是一个 batch 数据，包括图像信息和图像标签
zip 将图像信息和图像标签分开存储

batch[0]: 取出图像信息，将图像信息进行预处理

def collate_fn(batch):
  batch = list(zip(*batch))
  batch[0] = nested_tensor_from_tensor_list(batch[0])
  return tuple(batch)

nested_tensor_from_tensor_list 该函数有点类似 letterbox 函数的作用。

_max_by_axis: 找出一个 batch 图像中各个维度数的最大值，为了兼容不同图像大小
tensor_list: 一个 batch 的数据，维度为 [b, c, h, w]
tensor: 存放预处理后的数据，维度为 [b, c, h, w] ，先初始化为 0 ，然后将源图像的数据拷贝进去
mask: 维度为 [b, h, w] ，先初始化为 1 ，然后将原图像的位置设为 0
zip: 按 batch 维度进行提取数据，并设置 tensor 和 mask

NestedTensor: 该类对 tensor 和 mask 进行封装

def nested_tensor_from_tensor_list(tensor_list: List[Tensor]):
  # TODO make this more general
  if tensor_list[0].ndim == 3:
      if torchvision._is_tracing():
          # nested_tensor_from_tensor_list() does not export well to ONNX
          # call _onnx_nested_tensor_from_tensor_list() instead
          return _onnx_nested_tensor_from_tensor_list(tensor_list)

      # TODO make it support different-sized images
      max_size = _max_by_axis([list(img.shape) for img in tensor_list])
      # min_size = tuple(min(s) for s in zip(*[img.shape for img in tensor_list]))
      batch_shape = [len(tensor_list)] + max_size
      b, c, h, w = batch_shape
      dtype = tensor_list[0].dtype
      device = tensor_list[0].device
      tensor = torch.zeros(batch_shape, dtype=dtype, device=device)
      mask = torch.ones((b, h, w), dtype=torch.bool, device=device)
      for img, pad_img, m in zip(tensor_list, tensor, mask):
          pad_img[: img.shape[0], : img.shape[1], : img.shape[2]].copy_(img)
          m[: img.shape[1], :img.shape[2]] = False
  else:
      raise ValueError('not supported')
  return NestedTensor(tensor, mask)

NestedTensor 类比较简单，提供 decompose 函数获取 tensor 和 mask 。 repr 函数重定义了实例化对象的基本信息，调用 print(NestTensor) 时打印该函数的输出。

class NestedTensor(object):
    def __init__(self, tensors, mask: Optional[Tensor]):
        self.tensors = tensors
        self.mask = mask

    def to(self, device):
        # type: (Device) -> NestedTensor # noqa
        cast_tensor = self.tensors.to(device)
        mask = self.mask
        if mask is not None:
            assert mask is not None
            cast_mask = mask.to(device)
        else:
            cast_mask = None
        return NestedTensor(cast_tensor, cast_mask)

    def decompose(self):
        return self.tensors, self.mask

    def __repr__(self):
        return str(self.tensors)

2.1 backbone

DETR 网络中的 backbone 模块中包括两个部分，特征提取网络和位置编码。

arg.slr_backbone=1e-5 ，所以 train_backbone=True
return_interm_layers: 该变两用于语义分割
args.backbone: resnet50

args.dilation: False

def build_backbone(args):
  position_embedding = build_position_encoding(args)
  train_backbone = args.lr_backbone > 0
  return_interm_layers = args.masks
  backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)
  model = Joiner(backbone, position_embedding)
  model.num_channels = backbone.num_channels
  return model

我们先看特征提取网络，特征提取网络使用的是 resnet50 。

num_channels: resnet50 输出的特征图通道维度是 2048

class Backbone(BackboneBase):
  def __init__(self, name: str,
               train_backbone: bool,
               return_interm_layers: bool,
               dilation: bool):
      backbone = getattr(torchvision.models, name)(
          replace_stride_with_dilation=[False, False, dilation],
          pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d)
      num_channels = 512 if name in ('resnet18', 'resnet34') else 2048
      super().__init__(backbone, train_backbone, num_channels, return_interm_layers)

BackboneBase 是 Backbone 基类，该类提供了前向传播函数。

IntermediateLayerGetter: 该函数的作用是将 return_layers 之前的层保留，之后的层全部丢弃。
前向传播中， self.body 输出 OrderedDict 类型，保存特征图。
遍历 OrderedDict ，取出 mask ，对 mask 进行下采样 32 倍，因为 backbone 将图像下采样了 32 倍。
m[None]: [b,h,w]->[1,b,h,w] ，这么做的原因是 interpolate 只能接收 4 维的输入。
size: 是特征图的大小，也就是将 mask 重新 resize 到特征图的大小。

out: 将下采样的数据和 mask 保存到 NestedTensor 中。

class BackboneBase(nn.Module):

  def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool):
      super().__init__()
      for name, parameter in backbone.named_parameters():
          if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name:
              parameter.requires_grad_(False)
      if return_interm_layers:
          return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"}
      else:
          return_layers = {'layer4': "0"}
      self.body = IntermediateLayerGetter(backbone, return_layers=return_layers)
      self.num_channels = num_channels

  def forward(self, tensor_list: NestedTensor):
      xs = self.body(tensor_list.tensors)
      out: Dict[str, NestedTensor] = {}
      for name, x in xs.items():
          m = tensor_list.mask
          assert m is not None
          mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0]
          out[name] = NestedTensor(x, mask)
      return out

下面看下位置编码的实现。位置编码官方实现了两种，一种是固定位置编码，另一种是自学习位置编码，这里就介绍固定位置编码。

位置编码要考虑 x, y 两个方向，图像中任意一个点 (h, w) 有一个位置，这个位置编码长度为 256 ，前 128 维代表 h 的位置编码，后 128 维代表 w 的位置编码，把这两个 128 维的向量拼接起来就得到一个 256 维的向量，它代表 (h, w) 的位置编码。位置编码的计算公式如下图所示：

args.hidden_dim = 256

args.position_embedding=sine

def build_position_encoding(args):
  N_steps = args.hidden_dim // 2
  if args.position_embedding in ('v2', 'sine'):
      position_embedding = PositionEmbeddingSine(N_steps, normalize=True)
  return position_embedding

num_pos_feats=128
normalize=True
not_mask: mask 在之前已经将合法的图像位置设为了 0 ，取反后将合法图像位置设为了 1 。
y_embed, x_embed: 一开始计算 h, w 维度的累加和，然后除以累加后的和进行归一化处理。 cumsum 的功能是返回给定 axis 上的累加和
y_embed[:, -1:, :] 和 x_embed[:, :, -1:]: 取 y_embed 和 x_embed 累加后的最大值
eps: 为了防止除数为 0
dim_t: 先创建一个 1 维的长度维 128 的向量，根据位置编码的公式计算正余玄函数的参数 10000^(2i/128)
pos_x, pos_y: 先对 y_embed 和 x_embed 升维，维度变为 [b, h, w, 1] ，运用广播机制除以 dim_t 维度变为 [b, h, w, 128]
根据位置计算 sin 和 cos 值，使用 stack 进行拼接，拼接维度维 dim=4 ，注意数组的切片步长是 2 ，因此维度变为了 [b, h, w, 64, 2] ，最后使用 flatten(3) 再将维度变为 [b, h, w, 128]

pos: 最后使用 torch.cat 将 pos_x, pos_y 拼接在一起， dim=3 ，维度变为 [b, h, w, 256] ，然后使用 permute(0, 3, 1, 2) 将维度变为 [b, 256, h, w]

class PositionEmbeddingSine(nn.Module):
  def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
      super().__init__()
      self.num_pos_feats = num_pos_feats
      self.temperature = temperature
      self.normalize = normalize
      if scale is not None and normalize is False:
          raise ValueError("normalize should be True if scale is passed")
      if scale is None:
          scale = 2 * math.pi
      self.scale = scale

  def forward(self, tensor_list: NestedTensor):
      x = tensor_list.tensors
      # mask: [b, h, w]
      mask = tensor_list.mask
      assert mask is not None
      not_mask = ~mask
      y_embed = not_mask.cumsum(1, dtype=torch.float32)
      x_embed = not_mask.cumsum(2, dtype=torch.float32)
      if self.normalize:
          eps = 1e-6
          y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
          x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale
      dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
      dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)
      pos_x = x_embed[:, :, :, None] / dim_t
      pos_y = y_embed[:, :, :, None] / dim_t
      pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
      pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
      pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
      return pos

Joiner 类将特征提取网络和位置编码进行封装。

init 函数将特征提取网络和位置编码添加到 nn.Sequential 结构中。
self[0]: 指只特征提取网络
self[1]: 指 position_embedding
xs: 字典类型，值的类型为 NestedTensor ， BackboneBase 的前向传播的输出
self1: 对特征图进行位置编码
out: 链表，保存 NestedTensor 类型，维度为 [b, 2048, h, w]

pos: 链表，保存维度 [b, 256, h, w]

class Joiner(nn.Sequential):
  def __init__(self, backbone, position_embedding):
      super().__init__(backbone, position_embedding)

  def forward(self, tensor_list: NestedTensor):
      xs = self[0](tensor_list)
      out: List[NestedTensor] = []
      pos = []
      for name, x in xs.items():
          out.append(x)
          pos.append(self[1](x).to(x.tensors.dtype))
      return out, pos

总结一下 backbone ，首先使用 resnet50 网络对输入图像数据进行特征提取，数据维度由 [batch, 3, H, W] -> [batch, C, H1, W1] ，其中 C=2048 ， H1=H/32 W1=W/32 。然后根据特征图生成位置编码，维度为 [batch, 256, H1, W1] 。

2.2. transformer

DETR 中的 transformer ，包括 Encoder 和 Decoder 两个部分，网络结构如下图所示：

下面看 Transformer 的实现。

def build_transformer(args):
    return Transformer(
        d_model=args.hidden_dim,
        dropout=args.dropout,
        nhead=args.nheads,
        dim_feedforward=args.dim_feedforward,
        num_encoder_layers=args.enc_layers,
        num_decoder_layers=args.dec_layers,
        normalize_before=args.pre_norm,
        return_intermediate_dec=True,
    )

d_model=256
dim_feedforward=2048
normalize_before=False
encoder_norm=None

return_intermediate_dec=True: 是否保存 Decoder 的中间层用于计算损失， True 表示会保存。

class Transformer(nn.Module):
  def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
               num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
               activation="relu", normalize_before=False,
               return_intermediate_dec=False):
      super().__init__()
      encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                              dropout, activation, normalize_before)
      encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
      self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
      decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                              dropout, activation, normalize_before)
      decoder_norm = nn.LayerNorm(d_model)
      self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                        return_intermediate=return_intermediate_dec)
      self._reset_parameters()
      self.d_model = d_model
      self.nhead = nhead

  def _reset_parameters(self):
      for p in self.parameters():
          if p.dim() > 1:
              nn.init.xavier_uniform_(p)

先看 Encoder 。 Encoder 中包含多个 TransformerEncoderLayer 。

Encoder 由两部分组成，多头自注意力机制和 FFN
with_pos_embed: 该函数用来增加残差

DETR 网络的只有 q 和 k 需要增加位置编码

class TransformerEncoderLayer(nn.Module):

  def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
               activation="relu", normalize_before=False):
      super().__init__()
      self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
      self.linear1 = nn.Linear(d_model, dim_feedforward)
      self.dropout = nn.Dropout(dropout)
      self.linear2 = nn.Linear(dim_feedforward, d_model)
      self.norm1 = nn.LayerNorm(d_model)
      self.norm2 = nn.LayerNorm(d_model)
      self.dropout1 = nn.Dropout(dropout)
      self.dropout2 = nn.Dropout(dropout)
      self.activation = _get_activation_fn(activation)
      self.normalize_before = normalize_before

  def with_pos_embed(self, tensor, pos: Optional[Tensor]):
      return tensor if pos is None else tensor + pos

  def forward_post(self,
                   src,
                   src_mask: Optional[Tensor] = None,
                   src_key_padding_mask: Optional[Tensor] = None,
                   pos: Optional[Tensor] = None):
      # q 和 k 加上位置编码
      q = k = self.with_pos_embed(src, pos)
      # 多头注意力机制
      src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
                            key_padding_mask=src_key_padding_mask)[0]
      # 残差
      src = src + self.dropout1(src2)
      # 标准化
      src = self.norm1(src)
      # FFN: 第一个全链接接后跟 relu 激活函数，第二个全链接无激活函数
      src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
      # 残差
      src = src + self.dropout2(src2)
      # 边准化
      src = self.norm2(src)
      return src

  def forward(self, src,
              src_mask: Optional[Tensor] = None,
              src_key_padding_mask: Optional[Tensor] = None,
              pos: Optional[Tensor] = None):
      return self.forward_post(src, src_mask, src_key_padding_mask, pos)

TransformerEncoder 就是重复多次 TransformerEncoderLayer 。

src: [hxw, b, 256]
output: 第一次是 src ，之后都是上一个的输出
mask: None

src_key_padding_mask: [b, hxw]

class TransformerEncoder(nn.Module):

  def __init__(self, encoder_layer, num_layers, norm=None):
      super().__init__()
      self.layers = _get_clones(encoder_layer, num_layers)
      self.num_layers = num_layers
      # self.norm = None
      self.norm = norm

  def forward(self, src,
              mask: Optional[Tensor] = None,
              src_key_padding_mask: Optional[Tensor] = None,
              pos: Optional[Tensor] = None):
      output = src

      for layer in self.layers:
          output = layer(output, src_mask=mask,
                         src_key_padding_mask=src_key_padding_mask, pos=pos)

      if self.norm is not None:
          output = self.norm(output)

      return output

下面介绍 Decoder 。 Decoder 中包含多个 TransformerEncoderLayer 。

Decoder 由三部分构成，多头自注意力机制、多头注意力机制、 FNN
第一个多头自注意力机制的 q 和 k 加上 Object query

第二个多头注意力机制的 k 和 v 来自 Encoder ，且 k 加上了位置编码， q 来自第一个多头自注意力的输出，并且加上了 Object query

class TransformerDecoderLayer(nn.Module):

  def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1,
               activation="relu", normalize_before=False):
      super().__init__()
      self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
      self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
      self.linear1 = nn.Linear(d_model, dim_feedforward)
      self.dropout = nn.Dropout(dropout)
      self.linear2 = nn.Linear(dim_feedforward, d_model)

      self.norm1 = nn.LayerNorm(d_model)
      self.norm2 = nn.LayerNorm(d_model)
      self.norm3 = nn.LayerNorm(d_model)
      self.dropout1 = nn.Dropout(dropout)
      self.dropout2 = nn.Dropout(dropout)
      self.dropout3 = nn.Dropout(dropout)
      self.activation = _get_activation_fn(activation)
      self.normalize_before = normalize_before

  def with_pos_embed(self, tensor, pos: Optional[Tensor]):
      return tensor if pos is None else tensor + pos

  def forward_post(self, tgt, memory,
                   tgt_mask: Optional[Tensor] = None,
                   memory_mask: Optional[Tensor] = None,
                   tgt_key_padding_mask: Optional[Tensor] = None,
                   memory_key_padding_mask: Optional[Tensor] = None,
                   pos: Optional[Tensor] = None,
                   query_pos: Optional[Tensor] = None):
      # q 和 k 加上 Object query
      q = k = self.with_pos_embed(tgt, query_pos)
      # 多头自注意力机制
      tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
                            key_padding_mask=tgt_key_padding_mask)[0]
      # 残差
      tgt = tgt + self.dropout1(tgt2)
      # 标准化
      tgt = self.norm1(tgt)
      # 多头注意力机制，之前的 tgt 加上 Object query
      tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                                 key=self.with_pos_embed(memory, pos),
                                 value=memory, attn_mask=memory_mask,
                                 key_padding_mask=memory_key_padding_mask)[0]
      # 残差
      tgt = tgt + self.dropout2(tgt2)
      标准化
      tgt = self.norm2(tgt)
      # FFN
      tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
      # 残差
      tgt = tgt + self.dropout3(tgt2)
      # 标准化
      tgt = self.norm3(tgt)
      return tgt

  def forward(self, tgt, memory,
              tgt_mask: Optional[Tensor] = None,
              memory_mask: Optional[Tensor] = None,
              tgt_key_padding_mask: Optional[Tensor] = None,
              memory_key_padding_mask: Optional[Tensor] = None,
              pos: Optional[Tensor] = None,
              query_pos: Optional[Tensor] = None):
      return self.forward_post(tgt, memory, tgt_mask, memory_mask,
                               tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)

TransformerDecoder 就是重复多次 TransformerDecoderLayer 。

memory: 是 encoder 的输出
tgt: [100, 4, 256] 初始全是 0
memory_mask: None
tgt_key_padding_mask: None
memory_key_padding_mask: [b, hxw]
pos: 与 encoder 相同的位置编码

query_pos: [100, 4, 256]

class TransformerDecoder(nn.Module):

  def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
      super().__init__()
      self.layers = _get_clones(decoder_layer, num_layers)
      self.num_layers = num_layers
      self.norm = norm
      self.return_intermediate = return_intermediate

  def forward(self, tgt, memory,
              tgt_mask: Optional[Tensor] = None,
              memory_mask: Optional[Tensor] = None,
              tgt_key_padding_mask: Optional[Tensor] = None,
              memory_key_padding_mask: Optional[Tensor] = None,
              pos: Optional[Tensor] = None,
              query_pos: Optional[Tensor] = None):
      output = tgt

      intermediate = []

      for layer in self.layers:
          output = layer(output, memory, tgt_mask=tgt_mask,
                         memory_mask=memory_mask,
                         tgt_key_padding_mask=tgt_key_padding_mask,
                         memory_key_padding_mask=memory_key_padding_mask,
                         pos=pos, query_pos=query_pos)
          # self.return_intermediate=True
          # 保存 Decoder 中间层
          if self.return_intermediate:
              intermediate.append(self.norm(output))

      # self.norm is not None
      if self.norm is not None:
          output = self.norm(output)
          if self.return_intermediate:
              intermediate.pop()
              intermediate.append(output)

      if self.return_intermediate:
          return torch.stack(intermediate)

      return output.unsqueeze(0)

介绍了 Encoder 和 Decoder 后，在来看 Transformer 类的前向传播函数：

src: [b, 256, H/32, H/32]
mask: [b, h, w]
query_embed: Parameter ，可训练参数

pos_embed: 位置编码 [b, 256, h, w] [100, 256]

  def forward(self, src, mask, query_embed, pos_embed):
      # flatten NxCxHxW to HWxNxC
      bs, c, h, w = src.shape
      # src: [b, 256, h, w]->[b, 256, hxw]->[hxw, b, 256]
      src = src.flatten(2).permute(2, 0, 1)
      # pos_embed: [b, 256, h, w]->[b, 256, hxw]->[hxw, b, 256]
      pos_embed = pos_embed.flatten(2).permute(2, 0, 1)
      # query_embed: 是 embed.weight
      # [100, 256] -> [100, 1, 256]
      # repeat: [100, 1, 256]->[100, 4, 256]
      query_embed = query_embed.unsqueeze(1).repeat(1, bs, 1)
      # mask: [b, h, w]->[b, hxw]
      mask = mask.flatten(1)

      # tgt: [100, 4, 256]
      tgt = torch.zeros_like(query_embed)
      memory = self.encoder(src, src_key_padding_mask=mask, pos=pos_embed)
      hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                        pos=pos_embed, query_pos=query_embed)
      return hs.transpose(1, 2), memory.permute(1, 2, 0).view(bs, c, h, w)

总结一下，首先， backbone 对输入的每一个 batch 的图像进行通道维度的拉伸以及宽高维度的压缩；

位置编码要考虑 x, y 两个方向，图像中任意一个点 (h, w) 有一个位置，这个位置编码长度为 256 ，前 128 维代表 h 的位置编码，后 128 维代表 w 的位置编码，把这两个 128 维的向量拼接起来就得到一个 256 维的向量，它代表 (h, w) 的位置编码，另外，注意位置编码作用位置， DETR 在每个 Encoder 的输入都加入了位置编码，且只对 Query 和 Key 使用，即只与 Query 和 Key 相加，不与 Value 进行相加，对于 Decoder 的第二个多头注意力机制的 k 也加入了位置编码的信息；

Encoder 的输入和输出维度都是 [H1xW1, batch, 256] 。

Decoder 的输入由多个部分组成， Encoder 的输出、 object queries 、上一个 Decoder 的输出。 Decoder 的输入一开始初始化成维度为 [100, batch, 256] 维的全部元素都为 0 的张量，和 Object queries 加在一起之后充当第 1 个 multi-head self-attention 的 Query 和 Key 。第一个 multi-head self-attention 的 Value 为 Decoder 的输入，也就是全0的张量。每个 Decoder 的第 2 个 multi-head self-attention ，它的 Key 和 Value 来自 Encoder 的输出张量，维度为 [H1xW1, batch, 256] ，其中 Key 值还进行位置编码。 Query 值一部分来自第 1 个 Add and Norm 的输出，维度为 [100, batch, 256] 的张量，另一部分来自 Object queries ，充当可学习的位置编码。所以，第 2 个 multi-head self-attention 的 Key 和 Value 的维度为 [H1xW1, batch, 256] ，而 Query 的维度为 [100, batch, 256] 。

Decoder 的输出为 [batch, 100, 256] 。

object queries 维度为 [100, batch, 256] ，类型为 nn.Embedding 说明这个张量是学习得到的， Object queries 充当的其实是位置编码的作用，只不过它是可以学习的位置编码。

2.3. DETR 模型定义

DETR 类将之前介绍的 backbone 和 transformer 组装在一起。

self.input_proj: 对主干网络输出的特征图进行通道压缩 [b, 2048, H/32, H/32] -> [b, 256, H/32, H/32]
self.class_embed: 得到类别预测信息
self.bbox_embed: 得到边界框预测信息
features, pos: 都是链表，分别保存 NestedTensor 和位置编码

在 Decoder 网络中保存了每一个 Decoder 的输出，所以这里的 hs 的维度为 [6, b, 100, 256]

class DETR(nn.Module):
  def __init__(self, backbone, transformer, num_classes, num_queries, aux_loss=False):
      super().__init__()
      # num_queries=100
      self.num_queries = num_queries
      self.transformer = transformer
      hidden_dim = transformer.d_model
      # hidden_dim=256
      self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
      self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
      self.query_embed = nn.Embedding(num_queries, hidden_dim)
      self.input_proj = nn.Conv2d(backbone.num_channels, hidden_dim, kernel_size=1)
      self.backbone = backbone
      self.aux_loss = aux_loss

  def forward(self, samples: NestedTensor):
      if isinstance(samples, (list, torch.Tensor)):
          samples = nested_tensor_from_tensor_list(samples)
      # features: 链表，保存 NestedTensor 类型
      # pos: 链表，保存维度 [b, 256, h, w]
      features, pos = self.backbone(samples)

      src, mask = features[-1].decompose()
      assert mask is not None
      # transformer: 输出元组，分别为 Decoder 和 Encoder 的输出
      # torch.nn.Parameter 是继承自 torch.Tensor 的子类，其主要作用是作为 nn.Module 中的可训练参数使用。它与 torch.Tensor 的区别就是 nn.Parameter 会自动被认为是 module 的可训练参数，即加入到 parameter() 这个迭代器中去；而 module 中非 nn.Parameter() 的普通 tensor 是不在 parameter 中的。
      hs = self.transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]

      # hs:[b, 100, 256]
      # [b, 100, class+1]
      outputs_class = self.class_embed(hs)
      # [b, 100, 4]
      outputs_coord = self.bbox_embed(hs).sigmoid()
      out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
      # 是否计算 Decoder 层的损失
      if self.aux_loss:
          out['aux_outputs'] = self._set_aux_loss(outputs_class, outputs_coord)
      return out

  # 这个装饰器向编译器表明，应该忽略一个函数或方法，并用引发异常来替换它。这允许您在模型中保留与TorchScript不兼容的代码，同时仍然导出模型。
  # 遍历除最后一个 Decoder 的输出外的另外 5 个 Decoder 层的输出。
  @torch.jit.unused
  def _set_aux_loss(self, outputs_class, outputs_coord):
      return [{'pred_logits': a, 'pred_boxes': b}
              for a, b in zip(outputs_class[:-1], outputs_coord[:-1])]

3. 网络训练

在上一章网络结构中我们只到网络最终输出的结果是类别预测信息，维度为 [b, 100, class+1] ；边界框预测信息，维度为 [b, 100, 4] ，那么这些预测信息该如何与 GT 值进行匹配来计算网络损失，从而达到训练模型的效果呢？

作者在训练中引入了匈牙利算法来计算最优匹配，匈牙利算法是用来计算二分图的最优匹配问题，我们可以认为网络预测的结果对应二分图的左边， GT 值对应二分图的右边，根据二分图的定义，左边各个节点之间不能相连，右边各个节点之间不能相连，只能左边节点与右边节点相连。两边进行匹配的原则是什么呢？是使左边和右边最相近的两个节点进行相连，在 DETR 网络中也就是两边节点损失函数最小的进行匹配，匈牙利算法是计算每一个预测值与每一 GT 值都进行计算损失函数，然后找到相似度最接近的进行匹配，然后在反向传播过程中，更新网络参数使得匹配后的两个节点的损失逐渐变小。

def build(args):
    matcher = build_matcher(args)
    criterion = SetCriterion(num_classes, matcher=matcher, weight_dict=weight_dict,
                             eos_coef=args.eos_coef, losses=losses)

我们首先看下匹配的流程。

args.set_cost_class = 1
args.set_cost_bbox = 5

args.set_cost_giou = 2

def build_matcher(args):
  return HungarianMatcher(cost_class=args.set_cost_class, cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou)

负责预测值与 GT 值匹配的是 HungarianMatcher 类。

匈牙利算法匹配使用的代价参数由三个部分组成，类别损失、 L1 Loss 损失、 GIOU 损失，详细的计算过程请看下面代码的注释。

class HungarianMatcher(nn.Module):
  def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1):
      super().__init__()
      self.cost_class = cost_class
      self.cost_bbox = cost_bbox
      self.cost_giou = cost_giou
      assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, "all costs cant be 0"

  # output: 类别预测和边界框预测
  # target: 标签
  @torch.no_grad()
  def forward(self, outputs, targets):
      bs, num_queries = outputs["pred_logits"].shape[:2]

      # We flatten to compute the cost matrices in a batch
      # start_dim=0, end_dim=1
      # n 个预测器，每个预测器输出预测的类别概率
      out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
      # n 个预测器，每个预测器输出预测边界框
      out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]

      # Also concat the target labels and boxes
      # 目标中类别索引和边界框
      tgt_ids = torch.cat([v["labels"] for v in targets])
      # tgt_bbox: [N, 4]
      tgt_bbox = torch.cat([v["boxes"] for v in targets])

      # NLL: negative log likelihood loss  负对数似然损失
      # Compute the classification cost. Contrary to the loss, we don't use the NLL,
      # but approximate it in 1 - proba[target class].
      # The 1 is a constant that doesn't change the matching, it can be ommitted.
      # 从所有类别输出中，找出预测类别结果中，类别 ID 是 tgt_ids 的预测结果
      # 对于每个预测结果，把目前 gt 里面有的所有类别值提取出来，其余值不需要参与匹配
      # 行：取每一行；列：只取 tgt_ids 对应的 N 列
      # 匈牙利算法的作用是选出哪一个预测器使预测的结果的误差最小
      # 这样计算的目的是在所有预测器中，预测为某一类别时，哪一个预测器预测的结果误差最小，误差最小的就与标签进行匹配
      # 这只是匈牙利算法 cost 的一部分，假如某一张图片有 3 个目标，匈牙利算法要计算网络输出的 100 个预测结果中，预测成这三个目标时哪个预测结果使误差最小，这只是匈牙利算法 cost 的一部分， 下面的 L1 loss 和 iou 是同样的道理。
      # cost_class: [batch_size * num_queries, N]
      cost_class = -out_prob[:, tgt_ids]

      # Compute the L1 cost between boxes
      # x1 (Tensor) – input tensor of shape B×P×M .
      # x2 (Tensor) – input tensor of shape B×R×M .
      # output (Tensor) – will have shape B×P×R
      # out_bbox: [batch_size * num_queries, 4]
      # tgt_bbox: [N, 4]
      # cost_bbox: [batch_size * num_queries, N]
      cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)

      # Compute the giou cost betwen boxes
      # cost_giou: [batch_size * num_queries, N]
      cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))

      # Final cost matrix
      # [batch_size * num_queries, N]
      C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou
      # # [batch_size * num_queries, N] -> # [batch_size , num_queries, N]
      C = C.view(bs, num_queries, -1).cpu()

      # 计算一个 batch 中每一张图片中目标的大小
      sizes = [len(v["boxes"]) for v in targets]
      # torch.split(tensor, split_size_or_sections, dim=0)
      # 按照每张图像的 target 个数划分，计算每张图片的匹配情况
      # i 表示 batch 的索引
      # c[i]: 表示某一张图片中 100 个预测器与目标进行匹配的所有损失值，匈牙利算法要计算出最优的匹配，使损失值最小
      # indices: 是一个 tuple 类型，包含两个元素，第一个元素是匹配的行索引，第二个元素是匹配的列索引
      indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
      # 把匹配的行列索引转换成 tensor 类型，然后添加到列表里，每一项存储一个 tuple 类型， tuple 里有两个元素，匹配的行索引和列索引
      return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]

下面看 DETR 计算损失函数的类 SetCriterion 。

self.eos_coef: 背景类别的相对权重，这里是 0.1 ，这么设置的原因我的理解是背景的数目远远大于有物体的数目，所以计算损失时降低背景的损失的权重。
SetCriterion 先使用 HungarianMatcher 计算模型预测值与 gt 之间的匹配关系，然后对匹配后的结果计算损失。
损失的计算包括 label 损失，对应函数是 loss_labels ； boxes 损失，对应函数是 loss_boxes ； cardinality 损失，对应函数是 loss_cardinality ； cardinality 损失是计算预测有物体的个数的绝对损失，值是为了记录，不参与反向传播。
aux_outputs: 计算 Decoder 辅助损失，也就是前 5 个 Decoder 输出的损失

具体损失计算过程清参看代码中的注释。

# 
# 监测每一对匹配了的 gt 和预测值
class SetCriterion(nn.Module):
  def __init__(self, num_classes, matcher, weight_dict, eos_coef, losses):
      super().__init__()
      self.num_classes = num_classes
      self.matcher = matcher
      self.weight_dict = weight_dict
      self.eos_coef = eos_coef
      # losses = ['labels', 'boxes', 'cardinality']
      self.losses = losses
      empty_weight = torch.ones(self.num_classes + 1)
      empty_weight[-1] = self.eos_coef
      self.register_buffer('empty_weight', empty_weight)

  def loss_labels(self, outputs, targets, indices, num_boxes, log=True):
      """Classification loss (NLL)
      targets dicts must contain the key "labels" containing a tensor of dim [nb_target_boxes]
      """
      assert 'pred_logits' in outputs
      src_logits = outputs['pred_logits']

      # idx: batch_idx, src_idx: 记录每张图片中匹配成功的预测器索引
      idx = self._get_src_permutation_idx(indices)
      # target_classes_o: 保存每一张图片中，匹配到的类别索引
      target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(targets, indices)])
      # 初始化 target_classes 全为背景
      # [bs, 100]
      target_classes = torch.full(src_logits.shape[:2], self.num_classes,
                                  dtype=torch.int64, device=src_logits.device)
      # 设置匹配的预测器对应的类别
      target_classes[idx] = target_classes_o

      # 计算类别损失
      # src_logits: [bs, 100, 92] -> [bs, 92, 100] ，这里做维度变换的原因是 cross_entropy 里的 log_softmax 函数作用的维度是 dim=1 ，所以将预测类别维度放到 dim=1 位置
      # target_classes: [bs, 100]: 每个预测器匹配后的类别，大部分是背景
      loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
      losses = {'loss_ce': loss_ce}

      if log:
          # TODO this should probably be a separate loss, not hacked in this one here
          losses['class_error'] = 100 - accuracy(src_logits[idx], target_classes_o)[0]
      return losses

  @torch.no_grad()
  def loss_cardinality(self, outputs, targets, indices, num_boxes):
      """ Compute the cardinality error, ie the absolute error in the number of predicted non-empty boxes
      This is not really a loss, it is intended for logging purposes only. It doesn't propagate gradients
      """
      # 计算预测有物体的个数的绝对损失，值是为了记录，不参与反向传播
      pred_logits = outputs['pred_logits']
      device = pred_logits.device
      # tgt_lengths: 每张图像中目标的个数
      tgt_lengths = torch.as_tensor([len(v["labels"]) for v in targets], device=device)
      # Count the number of predictions that are NOT "no-object" (which is the last class)
      # 计算预测类别概率最大的索引不是背景的所有预测
      # pred_logits:[bs, 100, num_class]
      # (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1): [bs, 100]
      # card_pred: [bs], 保存每张图像中，预测有目标的个数
      card_pred = (pred_logits.argmax(-1) != pred_logits.shape[-1] - 1).sum(1)
      card_err = F.l1_loss(card_pred.float(), tgt_lengths.float())
      losses = {'cardinality_error': card_err}
      return losses

  def loss_boxes(self, outputs, targets, indices, num_boxes):
      """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
         targets dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
         The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size.
      """
      assert 'pred_boxes' in outputs
      idx = self._get_src_permutation_idx(indices)
      # outputs['pred_boxes']: [bs, 100, 4]
      # src_boxes: [N, 4]
      src_boxes = outputs['pred_boxes'][idx]
      target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(targets, indices)], dim=0)

      loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')

      losses = {}
      losses['loss_bbox'] = loss_bbox.sum() / num_boxes

      # generalized_box_iou: 两两计算 iou
      # torch.diag: 只取匹配后的 iou
      loss_giou = 1 - torch.diag(box_ops.generalized_box_iou(
          box_ops.box_cxcywh_to_xyxy(src_boxes),
          box_ops.box_cxcywh_to_xyxy(target_boxes)))
      losses['loss_giou'] = loss_giou.sum() / num_boxes
      return losses

  def loss_masks(self, outputs, targets, indices, num_boxes):
      """Compute the losses related to the masks: the focal loss and the dice loss.
         targets dicts must contain the key "masks" containing a tensor of dim [nb_target_boxes, h, w]
      """
      assert "pred_masks" in outputs

      src_idx = self._get_src_permutation_idx(indices)
      tgt_idx = self._get_tgt_permutation_idx(indices)
      src_masks = outputs["pred_masks"]
      src_masks = src_masks[src_idx]
      masks = [t["masks"] for t in targets]
      # TODO use valid to mask invalid areas due to padding in loss
      target_masks, valid = nested_tensor_from_tensor_list(masks).decompose()
      target_masks = target_masks.to(src_masks)
      target_masks = target_masks[tgt_idx]

      # upsample predictions to the target size
      src_masks = interpolate(src_masks[:, None], size=target_masks.shape[-2:],
                              mode="bilinear", align_corners=False)
      src_masks = src_masks[:, 0].flatten(1)

      target_masks = target_masks.flatten(1)
      target_masks = target_masks.view(src_masks.shape)
      losses = {
          "loss_mask": sigmoid_focal_loss(src_masks, target_masks, num_boxes),
          "loss_dice": dice_loss(src_masks, target_masks, num_boxes),
      }
      return losses

  def _get_src_permutation_idx(self, indices):
      # permute predictions following indices
      # indices 是一个列表，列表中的每一项表示每一张图片的预测值与真实值的匹配情况
      # src: 表示行索引，表示哪个预测器，形状为 (m, )
      # batch_idx: 记录每个目标在这个 batch 中的哪张图片，例如 [0,0,0,1,1,1,1,2,2,3]
      # 表示，图片 0 有三个目标，图片 2 有 4 个目标，以此类推。
      batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
      # src_idx: 记录每张图片中匹配成功的预测器索引，与 batch_idx 相对应
      src_idx = torch.cat([src for (src, _) in indices])
      return batch_idx, src_idx

  def _get_tgt_permutation_idx(self, indices):
      # permute targets following indices
      # 计算列索引
      batch_idx = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
      tgt_idx = torch.cat([tgt for (_, tgt) in indices])
      return batch_idx, tgt_idx

  # losses = ['labels', 'boxes', 'cardinality']
  def get_loss(self, loss, outputs, targets, indices, num_boxes, **kwargs):
      loss_map = {
          'labels': self.loss_labels,
          'cardinality': self.loss_cardinality,
          'boxes': self.loss_boxes,
          'masks': self.loss_masks
      }
      assert loss in loss_map, f'do you really want to compute {loss} loss?'
      return loss_map[loss](outputs, targets, indices, num_boxes, **kwargs)

  # outputs: 字典，包含类别预测 pred_logits ，边界框预测 pred_boxes 和 aux_outputs
  # targets: 标签，包含类别和边界框信息
  def forward(self, outputs, targets):
      """ This performs the loss computation.
      Parameters:
           outputs: dict of tensors, see the output specification of the model for the format
           targets: list of dicts, such that len(targets) == batch_size.
                    The expected keys in each dict depends on the losses applied, see each loss' doc
      """
      # 取出 pred_logits 和 pred_boxes
      outputs_without_aux = {k: v for k, v in outputs.items() if k != 'aux_outputs'}

      # Retrieve the matching between the outputs of the last layer and the targets
      # indices: 列表，每一项存储一个 tuple 类型，tuple里有两个元素，匹配的行索引和列索引
      indices = self.matcher(outputs_without_aux, targets)

      # Compute the average number of target boxes accross all nodes, for normalization purposes
      # 计算一个 batch 中目标的总和
      num_boxes = sum(len(t["labels"]) for t in targets)
      num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=next(iter(outputs.values())).device)
      if is_dist_avail_and_initialized():
          torch.distributed.all_reduce(num_boxes)
      num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()

      # Compute all the requested losses
      losses = {}
      # self.losses = ['labels', 'boxes', 'cardinality']
      for loss in self.losses:
          losses.update(self.get_loss(loss, outputs, targets, indices, num_boxes))

      # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
      if 'aux_outputs' in outputs:
          for i, aux_outputs in enumerate(outputs['aux_outputs']):
              indices = self.matcher(aux_outputs, targets)
              for loss in self.losses:
                  if loss == 'masks':
                      # Intermediate masks losses are too costly to compute, we ignore them.
                      continue
                  kwargs = {}
                  if loss == 'labels':
                      # Logging is enabled only for the last layer
                      kwargs = {'log': False}
                  l_dict = self.get_loss(loss, aux_outputs, targets, indices, num_boxes, **kwargs)
                  l_dict = {k + f'_{i}': v for k, v in l_dict.items()}
                  losses.update(l_dict)

      return losses

DETR 网络结构及实现就介绍到这。

可信计算之远程证明

Fri, 09 Apr 2021 00:00:00 +0000

0. 参考资料

remote-attestation-optiga-tpm

本文参考了英飞凌开源的远程证明代码的实现，并在此基础上进行了修改，这份开源的代码有三个分支， master 分支是远程证明实现相关的文档； device 分支是远程证明的客户端实现； server 分支是远程证明的服务端实现。除此之外还参考内核完整性度量模块相关的代码。

1. 概述

最近做了一些可信计算相关的工作，阶段性进展是搞清楚了远程证明的原理以及本地使用 tpm 模拟器实现了一个简单的远程证明方案。

可信计算是一个提出很多年的课题，这些年一直存在于理论研究中，网上能找到的资料除了 TCG (可信计算组织) 撰写的官方文档外，就剩下那些纯理论的论文，很少涉及工程实践。在 github 上也很少有相关的开源项目。本文将介绍可信计算中的一个小分支——远程证明。

2. 远程证明理论

2.1. 什么是远程证明

远程证明是指一个节点将自己平台的某些信息使用约定的格式和协议向另一个节点报告，使得另一节点能够获得这些信息，并判定该平台的可信状态。其目的是保证两个节点的身份和安全属性符合对一方的要求，其平台状态是可靠的。而在可信计算平台上的远程证明是指在平台上使用身份认证密钥 AIK 对当前存储的 PCR 值进行签名，然后报告给远程挑战者其平台的状态，是建立在可信度量、可信报告基础之上的技术。

2.2. 远程证明的原理

简单远程证明过程的完成要经历两个阶段：可信度量与可信报告。可信度量是平台组件完整性度量值的计算与存储，可信报告是将平台组件的度量值报告给外来用户，然后用户通过验证度量值来判断平台和应用程序是否被篡改。

可信度量是指通过一定的方法按步骤度量并报告平台的状态。从系统加电启动，一直到最后应用程序每一步都需要度量，整个启动序列都遵循先度量再执行的原则。当前阶段的代码负责度量下一阶段即将要执行的代码，然后再将度量值扩展到 PCR 中，这样一级信任一级，以此保证平台的可信，保证环境的安全。当然任何信任关系中总是存在某种基础性的假设，必然存在默认环节的信任关系的基石。在一个信任关系依次传递的链条中，源头在启动的过程中是被假设为安全可信的，不会受到度量。

在系统进行可信度量时，除了度量结果要在 PCR 中进行扩展之外，还将具体每一步的度量操作、中间状态和度量结果保存下来，可以作为系统可信度量的详细步骤进行参考，存储度量日志（Storage Measurement Log，SML）用于存储这些信息。由于 PCR 被认为是不可篡改的，并且它所保存的值可以通过 SML 重新计算出来，因此 SML 一般不需要安全保护。

3. 内核完整性度量

内核完整性度量模块 (IMA) 提供了远程证明所需的度量并扩展 PCR 和度量报告。

内核完整性度量的实现在 security/integrity/ima 目录下。要想使用 IMA ，需要首先配置内核的 config 文件，使能 IMA 。还要在 config 文件中指定完整性度量扩展的 PCR 寄存器的索引值。

CONFIG_IMA=y
CONFIG_IMA_MEASURE_PCR_IDX=11

3.1. IMA 策略配置

策略配置在 ima_policy.c 文件中。内核默认不进行完整性度量，需要进行配置，可以在内核启动项中增加 ima_tcb 选项即可开启，这是内核的默认策略，该策略度量的范围比较大，我增加了一个只度量二进制程序的策略，实现如下：

在 ima_policy.c 文件中增加如下代码，配置度量二进制程序的策略。

static struct ima_rule_entry bprm_measurement_rules[] __ro_after_init = {
	{.action = MEASURE, .func = BPRM_CHECK, .mask = MAY_EXEC,
	 .flags = IMA_FUNC | IMA_MASK},
};

static int __init bprm_measure_policy_setup(char *str)
{
	if (bprm_policy)
		return 1;

	bprm_policy = 1;
	return 1;
}
__setup("bprm_tcb", bprm_measure_policy_setup);

在 ima_init_policy 函数中增加将策略添加到策略列表的代码。

void __init ima_init_policy(void)
{
    if (bprm_policy) {
		for (i = 0; i < ARRAY_SIZE(bprm_measurement_rules); i++)
			list_add_tail(&bprm_measurement_rules[i].list,
				      &ima_default_rules);
	}
}

这样在内核启动项中增加 bprm_tcb 启动项即可开启动量二进制程序的策略了。

3.2. IMA 模版

IMA 模版主要用来配置度量报告，配置度量报告要存哪些值，以及导出度量报告。

内核默认支持的模板如下，其中 name 在 Kconfig 中默认选择的是 ima-ng ，因此 fmt 选择的是 d-ng|n-ng 。

static struct ima_template_desc builtin_templates[] = {
	{.name = IMA_TEMPLATE_IMA_NAME, .fmt = IMA_TEMPLATE_IMA_FMT},
	{.name = "ima-ng", .fmt = "d-ng|n-ng"},
	{.name = "ima-sig", .fmt = "d-ng|n-ng|sig"},
	{.name = "", .fmt = ""},	/* placeholder for a custom format */
};

根据模版的 fmt ，从 supported_fields 中选出度量报告中存储文件 hash 值和文件名称。其中 field_init 函数用来生成度量报告， field_show 函数用来导出度量报告。

ima_eventdigest_ng_init: 存储数据格式为 [] + ':' + '\0' + digest ，以及数据长度
ima_eventname_ng_init: 存储数据为文件路径，以及路径长度
ima_show_template_digest_ng: 将该 field_id 对应的报告以长度、内容的格式导出，根据不同的 show 字段显示不同的格式

ima_show_template_string: 将该 field_id 对应的报告以长度、内容的格式导出，根据不同的 show 字段显示不同的格式

static struct ima_template_field supported_fields[] = {
  {.field_id = "d", .field_init = ima_eventdigest_init,
   .field_show = ima_show_template_digest},
  {.field_id = "n", .field_init = ima_eventname_init,
   .field_show = ima_show_template_string},
  {.field_id = "d-ng", .field_init = ima_eventdigest_ng_init,
   .field_show = ima_show_template_digest_ng},
  {.field_id = "n-ng", .field_init = ima_eventname_ng_init,
   .field_show = ima_show_template_string},
  {.field_id = "sig", .field_init = ima_eventsig_init,
   .field_show = ima_show_template_sig},
};

3.3. IMA 生成度量报告

IMA 根据策略来生成度量报告，比如我配置的策略是度量二进制，那么在有二进制运行时就进行度量并生成度量报告。

除了配置策略之外， IMA 有一个默认策略，那就是 boot_aggregate ，这是度量报告的第一条数据。我们以它为例介绍 IMA 生成度量报告的流程。

度量日志的第一条记录是 boot_aggregate ，如果存在 tpm 芯片的话它的 hash 值是索引为 0~7 的 PCR 的值进行 hash ，在 ima_calc_boot_aggregate 函数中实现；如果没有 tpm 芯片的话则设为 0 。
ima_alloc_init_template: 该函数根据模版进行初始化 entry ，主要填充 struct ima_field_data 字段。

ima_store_template: 该函数将初始化后的 entry 中 struct ima_field_data 字段的数据计算 hash 值，保存到 digest 字段中，然后将 entry 添加到全局链表中，如果 tpm 芯片存在的话，将 digest 字段的值扩展到 PCR 中。

static int __init ima_add_boot_aggregate(void)
{
  static const char op[] = "add_boot_aggregate";
  const char *audit_cause = "ENOMEM";
  struct ima_template_entry *entry;
  struct integrity_iint_cache tmp_iint, *iint = &tmp_iint;
  struct ima_event_data event_data = {iint, NULL, boot_aggregate_name,
                      NULL, 0, NULL};
  int result = -ENOMEM;
  int violation = 0;
  struct {
      struct ima_digest_data hdr;
      char digest[TPM_DIGEST_SIZE];
  } hash;

  memset(iint, 0, sizeof(*iint));
  memset(&hash, 0, sizeof(hash));
  iint->ima_hash = &hash.hdr;
  iint->ima_hash->algo = HASH_ALGO_SHA1;
  iint->ima_hash->length = SHA1_DIGEST_SIZE;

  if (ima_tpm_chip) {
      result = ima_calc_boot_aggregate(&hash.hdr);
      if (result < 0) {
          audit_cause = "hashing_error";
          goto err_out;
      }
  }

  result = ima_alloc_init_template(&event_data, &entry);
  if (result < 0) {
      audit_cause = "alloc_entry";
      goto err_out;
  }

  result = ima_store_template(entry, violation, NULL,
                  boot_aggregate_name,
                  CONFIG_IMA_MEASURE_PCR_IDX);
  if (result < 0) {
      ima_free_template_entry(entry);
      audit_cause = "store_entry";
      goto err_out;
  }
  return 0;
err_out:
  integrity_audit_msg(AUDIT_INTEGRITY_PCR, NULL, boot_aggregate_name, op,
              audit_cause, result, 0);
  return result;
}

创建 entry ，并调用 field_init 函数对其初始化

int ima_alloc_init_template(struct ima_event_data *event_data,
              struct ima_template_entry **entry)
{
  struct ima_template_desc *template_desc = ima_template_desc_current();
  int i, result = 0;

  *entry = kzalloc(sizeof(**entry) + template_desc->num_fields *
           sizeof(struct ima_field_data), GFP_NOFS);
  if (!*entry)
      return -ENOMEM;

  (*entry)->template_desc = template_desc;
  for (i = 0; i < template_desc->num_fields; i++) {
      struct ima_template_field *field = template_desc->fields[i];
      u32 len;

      result = field->field_init(event_data,
                     &((*entry)->template_data[i]));
      if (result != 0)
          goto out;

      len = (*entry)->template_data[i].len;
      (*entry)->template_data_len += sizeof(len);
      (*entry)->template_data_len += len;
  }
  return 0;
out:
  ima_free_template_entry(*entry);
  *entry = NULL;
  return result;
}

ima_calc_field_array_hash: 对所有的 field 数据进行 hash 计算，包括长度，将 hash 值保存在 digest 字段

ima_add_template_entry: 将 entry 添加到全局链表，导出度量报告时需要使用。最后将 hash 值扩展到 PCR 中。

int ima_store_template(struct ima_template_entry *entry,
             int violation, struct inode *inode,
             const unsigned char *filename, int pcr)
{
  static const char op[] = "add_template_measure";
  static const char audit_cause[] = "hashing_error";
  char *template_name = entry->template_desc->name;
  int result;
  struct {
      struct ima_digest_data hdr;
      char digest[TPM_DIGEST_SIZE];
  } hash;

  if (!violation) {
      int num_fields = entry->template_desc->num_fields;

      /* this function uses default algo */
      hash.hdr.algo = HASH_ALGO_SHA1;
      result = ima_calc_field_array_hash(&entry->template_data[0],
                         entry->template_desc,
                         num_fields, &hash.hdr);
      if (result < 0) {
          integrity_audit_msg(AUDIT_INTEGRITY_PCR, inode,
                      template_name, op,
                      audit_cause, result, 0);
          return result;
      }
      memcpy(entry->digest, hash.hdr.digest, hash.hdr.length);
  }
  entry->pcr = pcr;
  result = ima_add_template_entry(entry, violation, op, inode, filename);
  return result;
}

总结一下， boot_aggregate 的度量报告生成流程：

计算索引 0~7 的 PCR 值的hash值，如果没有 tpm 芯片 hash 值为 0
将 [] + ':' + '\0' + digest 和数据长度写入 struct ima_field_data 结构中
将 ‘boot_aggregate’ + ‘\0’ 和数据长度写入struct ima_field_data 结构中
将上述两步进行 hash 处理，保存到 digest 字段中
将 entry 添加到全局链表，如果有 tpm 芯片将 digest 值扩展到 PCR 11 中

3.4. IMA 导出度量报告

在进行远程证明时需要将度量报告导出发给远端服务器，内核帮我们提供了导出度量报告的接口，在 ima_fs.c 文件中。

导出度量报告有两个接口，一个是导出 ascii 格式的，一个是导出二进制格式的，其中 ascii 码格式是可视化的，这两种格式都可以作为远程证明的度量报告，在英飞凌的源码中使用的是二进制格式的度量报告。
导出度量报告就是遍历全局的 entry 链表，将数据一条一条读取出来，内核在 /sys/kernel/security/ima 目录下提供了 ascii_runtime_measurements 和 binary_runtime_measurements 两个文件，直接读取即可。

最终是调用 field_show 函数，输出报告内容

int __init ima_fs_init(void)
{
  ascii_runtime_measurements =
      securityfs_create_file("ascii_runtime_measurements",
                 S_IRUSR | S_IRGRP, ima_dir, NULL,
                 &ima_ascii_measurements_ops);
  if (IS_ERR(ascii_runtime_measurements))
      goto out;

  runtime_measurements_count =
      securityfs_create_file("runtime_measurements_count",
                 S_IRUSR | S_IRGRP, ima_dir, NULL,
                 &ima_measurements_count_ops);
  if (IS_ERR(runtime_measurements_count))
      goto out;
}

4. tpm 模拟器环境搭建

首先安装 tpm 相关的工具包，包括 tpm2-tss-2.0 、 tpm2-abrmd-2.0 、 tpm2-tools-2.0 三个包，这里注意，我用的这几个包的版本比较老了，大部分命令与英飞凌的远程证明实现中不兼容，但命令的原理是相同的。

首先，我们要添加 /usr/lib/systemd/system/tpm2-abrmd.service 文件，这里要注意，使用 tpm 模拟器的话，要添加 ‘–tcti=libtss2-tcti-mssim.so.0:host=127.0.0.1,port=2321’ 选项。

[Unit]
Description=TPM2 Access Broker and Resource Management Daemon

[Service]
Type=dbus
BusName=com.intel.tss2.Tabrmd
StandardOutput=syslog
ExecStart=/usr/sbin/tpm2-abrmd --tcti=libtss2-tcti-mssim.so.0:host=127.0.0.1,port=2321
User=tss

[Install]
WantedBy=multi-user.target

还要添加 /etc/dbus-1/system.d/tpm2-abrmd.conf 文件，内容如下：

<!DOCTYPE busconfig PUBLIC "-//freedesktop//DTD D-BUS Bus Configuration 1.0//EN"
 "http://www.freedesktop.org/standards/dbus/1.0/busconfig.dtd">
<busconfig>
  <!-- ../system.conf have denied everything, so we just punch some holes -->
  <policy user="tss">
    <allow own="com.intel.tss2.Tabrmd"/>
  </policy>
  <policy user="root">
    <allow own="com.intel.tss2.Tabrmd"/>
  </policy>
  <policy context="default">
    <allow send_destination="com.intel.tss2.Tabrmd"/>
    <allow receive_sender="com.intel.tss2.Tabrmd"/>
  </policy>
</busconfig>

去 tpm官网下载 TPM 模拟器，编译安装，运行 tpm 模拟器 tpm_server -rm & ，这里的 -rm 作用是清空 tpm 模拟器的数据。

编译安装 tpm2-tss-2.0 和 tpm2-tools-2.0 。

然后启动 tpm2-abrmd ，执行 systemctl start tpm2-abrmd.service 。

最后执行 tpm2_pcrlist 检测 tpm模拟器是否能够正常运行。

5. 远程证明服务端

远程证明服务端是使用 SpringBoot 开发，项目由 maven 构建，所以需要安装 mvn 工具，并且代码中使用了 java-9 的特性，还需要安装 OpenJDK9 才能编译通过。具体流程如下：

yum install maven
tar -zxf OpenJDK9-OPENJ9_x64_Linux_jdk-9.0.4.12_openj9-0.9.0.tar.gz -C /usr/lib/jvm
alternatives --install /usr/bin/java java /usr/lib/jvm/jdk-9.0.4+12/bin/java 2
alternatives --config java
export JAVA_HOME=/usr/lib/jvm/jdk-9.0.4+12
mvn install
sudo java -jar server-0.0.1-SNAPSHOT.jar

由于我使用的 tpm 模拟器，tpm 命令无法操作 nv 。而在英飞凌的实现中，通过下面命令获取 ek 的证书。

# ek 证书默认存储在 0x1c00002 地址。
tpm2_nvread 0x1c00002 -o ek.crt

在我的实验中没有使用 ek 证书，直接让客户端将 ek 的公钥发送到服务段。所以证明的过程中没有实现 ek 证书有效性的认证。

另外一个修改点是 ek 公钥的长度大于了 600 个字节，所以在 User.java 中将长度改为 1000 。还有就是我内核生成度量日志使用的模版是 ima-ng ，而英飞凌的实现是按照 ima-sig 格式解析的，两者差异是 ima-sig 多了一个签名值，所以要在代码里将签名值的解析屏蔽掉。

6. 远程证明客户端

英飞凌的实现中有一个 bug ， fByteAry2HexStr 函数的第一个参数定义成 char * 类型， char 类型最高位为符号位，如果符号位为 1 时，转成 hexstr 时为 0xff ，这里导致了传输公钥时出错，而公钥中存在着大量超过 128 的数据。

另外，就是英飞凌实现的 tpm 指令与我的环境不兼容。

初始化 tpm 模拟器指令：

tpm2_takeownership -c

tpm2_takeownership -o ownerpasswd -e endorsepasswd -l lockpasswd

tpm2_getpubek -o ownerpasswd -e endorsepasswd -g rsa -f ek.pub -H 0x81010001

tpm2_getpubak -o ownerpasswd -e endorsepasswd -E 0x81010001 -g rsa -D sha256 -s rsassa -k 0x81010002 -f ak.pub -n ak.name

伪造 pcr 值，因为时 tpm 模拟器，所以要根据度量日志伪造 pcr 值：

#! /usr/bin/python

import os

with open("ascii_runtime_measurements", "r") as f:
    lines = f.readlines()
    for line in lines:
        pcr = line.split(" ")[1]
        os.system("tpm2_pcrextend 11:sha1=%s" % pcr)

执行远程证明指令：

tpm2_readpublic -H 0x81010002 -o ak.pub
tpm2_readpublic -H 0x81010001 -o ek.pub

tpm2_pcrlist -L sha1:0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18,19,20,21,22,23 -o pcr

./bin/attune
./bin/atelic
sh activecredential.sh
sh quote.sh
./bin/attest

基于 tpm 模拟器的远程证明实现就介绍到这里，详细代码实现可以查看 remote-attestation-optiga-tpm 代码。

transformer 在 CV 中的应用(一) Transformer 介绍

Mon, 22 Mar 2021 00:00:00 +0000

0. 参考资料

1. 概述

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型，经过近几年的发展， Transformer 不仅在 NLP 领域有很好的应用，在 CV 领域也得到了快速发展。最初 Transformer 是为了替换 RNN 网络而出现的，因为 RNN 循环神经网络是顺序模型，时间 t 时刻的计算依赖于时间 t-1 时刻的输出，这限制了模型的并行能力；其次顺序计算的过程中信息会丢失。为了解决上述两个问题， Transformer 使用了 attention 机制，将序列中任意两个位置之间的距离缩小为一个常量。；其次它避免了使用顺序结构，因此具有更好的并行性。

2. Transformer

如果把 Transformer 看成一个黑盒的话，那么 Transformer 的结构应该是如下图所示，输入一个句子，输出要翻译的结果。

但实际上，这个黑盒的内部是由一些列 Encode-Decoder 结构组成的，如下图所示：

论文中给出了编码器和解码器的个数是 6 个，如下图所示：

2.1. 自注意力机制

自注意力机制是 Transformer 最核心的部分。在介绍自注意力机制之前，我们有必要先来了解一下注意力机制。

如上图所示，在翻译任务中，是将 Source 语言翻译成 Target 语言，将 Source 中的构成元素想象成是由一系列的 <Key,Value> 数据对构成，此时给定 Target 中的某个元素 Query ，通过计算 Query 和各个 Key 的相似性或者相关性，得到每个 Key 对应 Value 的权重系数，然后对 Value 进行加权求和，即得到了最终的 Attention 数值。所以本质上 Attention 机制是对 Source 中元素的 Value 值进行加权求和，而 Query 和 Key 用来计算对应 Value 的权重系数。即可以将其本质思想改写为如下公式：

自注意力机制主要是用来表示句子内的关系，比如句子中有一个单词 ‘it’ ，它指代的上下文中的什么东西。在自注意力机制中有 3 个向量，分别是 Query 、 Key 、 Value ，长度都相同，假如是 512 。他们是输入通过点乘 3 个不同的权值矩阵 W 得到，这三个权值矩阵的尺寸都是 512x512 。这三个权值矩阵是网络通过反向传播学习到的。

Query 、 Key 、 Value 矩阵的计算示例图如下所示：

如上图所示，假如输入两个单词，每个单词的词向量的维度为 [1, 512] ，那么两个单词组成的词向量矩阵维度为 [2, 512] ，用词向量分别乘以三个维度为 [512, 512] 的权值矩阵，就得到维度为 [2, 512] 的三个矩阵，分别是 Query 、 Key 、 Value 矩阵。

得到 Query(q) 、 Key(k) 、 Value(v) 矩阵后，为每个向量计算 Score ， Score = qk.T 。得到的 Score 的维度为 [2, 2] 。然后对 Score 除以 Query 矩阵维度的平方根后进行 softmax 激活，得到维度为 [2, 2] 的矩阵，最后用该维度为 [2, 2] 矩阵点乘 v 矩阵，得到矩阵 Z ，维度为 [2, 512] 。上述 Score 除以 Query 矩阵维度的平方根是为了防止 qk.T 的数值会随着维度的增大而增大，所以要除以该值，相当于归一化的效果。具体流程如下图所示：

Query 、 Key 、 Value 矩阵的计算流程如下图所示：

矩阵 Z 的计算流程如下图所示：

这里还要提一下带 masked 的注意力机制，它在 Transformer 的 Decoder 模块中使用， masked 的意思是使 attention 只会关注已经产生的序列，防止预测时不会受到未来的信息干扰。

了解了自注意力机制是如何工作的，再看自注意力机制的输出计算公式就很好理解了，如下图所示：

下面看自注意力机制的代码实现。我们首先看下自注意力机制的框架图，如下图所示，自注意力机制的实现代码就是对这张图的翻译过程。

self.temperature: 起到归一化作用
k.transpose: 将 k 变成 k.T ，由于数据是一个 batch 输入的，所以用 transpose 将 2 和 3 个维度进行交换
mask: 带 masked 的注意力机制。

masked_fill: 将 mask 数组中为 0 位置的 attn 值设为 -1e9 ，作用是将 mask 为 0 的位置值屏蔽

class ScaledDotProductAttention(nn.Module):
  ''' Scaled Dot-Product Attention '''

  def __init__(self, temperature, attn_dropout=0.1):
      super().__init__()
      self.temperature = temperature
      self.dropout = nn.Dropout(attn_dropout)

  def forward(self, q, k, v, mask=None):
      attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
      if mask is not None:
          attn = attn.masked_fill(mask == 0, -1e9)
      attn = self.dropout(F.softmax(attn, dim=-1))
      output = torch.matmul(attn, v)
      return output, attn

Query，Key，Value 的概念取自于信息检索系统，当你在某电商平台搜索某件商品时，你在搜索引擎上输入的内容便是 Query ，然后搜索引擎根据 Query 为你匹配 Key （例如商品的种类，颜色，描述等），然后根据 Query 和 Key 的相似度得到匹配的内容（Value)。

self-attention 中的 Q，K，V 也是起着类似的作用，在矩阵计算中，点积是计算两个矩阵相似度的方法之一，因此 qk.T 进行相似度的计算。接着便是根据相似度进行输出的匹配，这里使用了加权匹配的方式，而权值就是 query 与 key 的相似度。

2.2. 多头自注意力机制

多头自注意力机制，简单理解是词与词之间的关系的产生可能是多种多样的。假如 head=8 ，则实际上使用 8 组不同的权值矩阵 W 分别计算 8 组 Query 、 Key 、 Value 矩阵，最后得到 8 组不同的 Z 值，将他们在列维度上进行拼接，最后将得到新的特征矩阵送入全链接层得到最终的输出 Z 。多头注意力机制实际上时将权值矩阵由 [512, 512] 分成 [512, 8*64] ，最终计算的维度与自注意力机制计算的结果维度是一样的。计算过程如下图所示：

下面我们看一下多头注意力机制的代码实现：

n_head: 多头注意力，默认是 8 头
d_k, d_v: 转换矩阵的维度，默认是 64 维。 q 矩阵和 k 矩阵维度相同
d_model: 输入是由 Word2Vec 等词嵌入方法将输入语料转化成特征向量，向量维度默认是 512
self.w_qs self.w_ks self.w_vs: 输入的特征向量分别乘上三个不同的转换矩阵得到 embedding 。
self.fc: 使用转换矩阵调整维度，将维度调整到 d_model ，这里的输出就是前面介绍的 Z 矩阵。
temperature: 值为 d_k ** 0.5 ，因为 qk.T 的数值会随着 dimension 的增大而增大，所以要除以 dimension 的平方根，相当于归一化的效果。
self.attention: 上文提到的注意力机制。

self.layer_norm: Layer Normalization 后面介绍。

class MultiHeadAttention(nn.Module):
  def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1):
      super().__init__()
      self.n_head = n_head
      self.d_k = d_k
      self.d_v = d_v
      self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False)
      self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False)
      self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False)
      self.fc = nn.Linear(n_head * d_v, d_model, bias=False)
      self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)
      self.dropout = nn.Dropout(dropout)
      self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

前向传播函数:

residual: 将输入保存，用来计算残差。
计算三个权值矩阵， q k v 。NLP 模型输入的数据维度为 [batch_size, seq_len, input_dim] ，其中， seq_len 表示 batch 中每一个序列的长度， input_dim 表示每个序列中的每一个单词的 embedding ，这里的 input_dim 就等于 d_model 。进入全链接层后，输出的维度为 [batch_size, seq_len, output_dim] ，这里的 output_dim 就是 n_head * d_k 。这里又将维度进行了转换，变为 [batch_size, seq_len, n_head, d_k] 。
通过 transpose 将 q k v 维度变为 [batch_size, n_head, seq_len, d_k] ，矩阵 q 和 v 施加注意力机制后得到 Z 矩阵形状为 [batch_size, n_head, seq_len, seq_len] ，得到的矩阵 Z 再与 v 进行点乘，形状变为 [batch_size, n_head, seq_len, d_k]
通过注意力机制计算 Z ，这里的 Z 存储在 q 中。
将 Z 值的维度在变为输入时的维度，以便进行残差计算，维度为 [batch_size, seq_len, n_head * d_k]
使用 self.fc 将维度调整到 d_model

最后将计算残差后的值进行 LN 计算。

def forward(self, q, k, v, mask=None):
  d_k, d_v, n_head = self.d_k, self.d_v, self.n_head
  sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
  residual = q
  q = self.w_qs(q).view(sz_b, len_q, n_head, d_k)
  k = self.w_ks(k).view(sz_b, len_k, n_head, d_k)
  v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)

  # Transpose for attention dot product: b x n x lq x dv
  q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

  if mask is not None:
      mask = mask.unsqueeze(1)   # For head axis broadcasting.

  q, attn = self.attention(q, k, v, mask=mask)

  q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
  q = self.dropout(self.fc(q))
  q += residual
  q = self.layer_norm(q)
  return q, attn

2.4. 位置编码

我们知道像 RNN 这种序列模型，天生就具有位置信息，但是 Transformer 是如何提供这种能力的呢？答案是位置编码。也就是在词向量中加入位置信息。那么如何增加位置信息呢？论文中提到有两种方式，一种是根据数据学习，另一种是手动设计编码规则，论文中作者采用了第二种方式。通常位置编码是一个长度为 d 维的特征向量，维度与词向量相同，这样便于和词向量进行相加操作。位置编码的计算公式如下图所示：

上式中 pos 表示单词的位置， i 表示单词的维度。准确来说 2i 和 2i + 1 表示单词的维度， i 的取值范围是 [0, d/2] 。作者提到，这样设计是因为 NLP 任务中，除了单词的绝对位置，单词的相对位置也非常重要。根据下图公式可知，任意位置 p+k 都可以被位置 k 的线性函数表示，这为模型捕捉单词之间的相对位置关系提供了非常大的便利。

位置编码代码实现：

n_position: 表示 token 在 sequence 中的位置

d_hid: 表示了 Positional Encoding 的维度

class PositionalEncoding(nn.Module):

  def __init__(self, d_hid, n_position=200):
      super(PositionalEncoding, self).__init__()

      self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))

  def _get_sinusoid_encoding_table(self, n_position, d_hid):
      ''' Sinusoid position encoding table '''
      # TODO: make it with torch instead of numpy

      def get_position_angle_vec(position):
          return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]

      sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)])
      sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2])  # dim 2i
      sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2])  # dim 2i+1

      return torch.FloatTensor(sinusoid_table).unsqueeze(0)

  def forward(self, x):
      return x + self.pos_table[:, :x.size(1)].clone().detach()

2.5. FFN

FFN 即是 Feed Forward Neural Network 的简称，其实就是两个全链接层。第一个全链接层使用 relu 激活函数，第二个全链接层使用线性激活函数。第二个全链接层的输出经过 dropout 后与输入进行残差计算，最后使用 LN 进行标准化。代码如下：

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_in, d_hid, dropout=0.1):
        super().__init__()
        self.w_1 = nn.Linear(d_in, d_hid)
        self.w_2 = nn.Linear(d_hid, d_in)
        self.layer_norm = nn.LayerNorm(d_in, eps=1e-6)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        residual = x
        x = self.w_2(F.relu(self.w_1(x)))
        x = self.dropout(x)
        x += residual
        x = self.layer_norm(x)
        return x

2.7. Layer Normolization

BN 是取不同样本的同一个通道的特征做归一化， BN 是按照样本数计算归一化统计量的，当样本数很少时，样本的均值和方差不能反映全局的统计分布息，所以基于少量样本的 BN 的效果会变得很差，在一些场景中，比如说硬件资源受限，在线学习等场景， BN 是非常不适用的； LN 则是取的是同一个样本的不同通道做归一化，即根据样本的特征数做归一化。 LayerNorm 中不会像 BatchNorm 那样跟踪统计全局的均值方差，因此 train() 和 eval() 对 LayerNorm 没有影响。 LN 和 BN 的区别可以看下图：

LayerNorm 有三个参数，含义分别是：

normalized_shape: 输入尺寸
eps: 归一化时加在分母上防止除零
elementwise_affine: 如果设为 False ，则 LayerNorm 层不含有任何可学习参数。如果设为 True (默认是 True) 则会包含可学习参数 weight 和 bias ，用于仿射变换，即对输入数据归一化到均值 0 方差 1 后，乘以 weight ，加上 bias 。

2.8. Encoder-Decoder 模块

了解了上述各个子模块的原理之后， Transformer 整体结构也就掌握了，先贴一张 Transformer 整体框架图，如下图所示，图中左侧方框框起来的是 Encoder ，右侧方框框起来的是 Decoder 。

这里主要分析下网络的输入和输出，首先看框架图的左侧 Encoder 模块， Encoder 模块的输入就是词向量与位置编码的和。 Encoder 的输出当作 Decoder 输入的一部分进入 Decoder 模块。

Decoder 的输入包括 2 部分，一部分来自下方，是前一个 timestep 的输出，再加上一个表示位置的 Positional Encoding ；另一部分来自 Encoder 的输出，作为中间的 attention 的 key 和 value ，而中间的 attention 的 query 来自第一个 attention 的输出。 Decoder 的输出是对应 i 位置的输出词的概率分布。 Decoder 的解码不是一次把所有序列解出来的，而是像 RNN 一样一个一个解出来的，因为要用上一个位置的输入当作 attention 的 query 。

Encoder 代码实现如下：

class EncoderLayer(nn.Module):
    ''' Compose with two layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(EncoderLayer, self).__init__()
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(self, enc_input, slf_attn_mask=None):
        enc_output, enc_slf_attn = self.slf_attn(
            enc_input, enc_input, enc_input, mask=slf_attn_mask)
        enc_output = self.pos_ffn(enc_output)
        return enc_output, enc_slf_attn

Decoder 代码实现如下：

class DecoderLayer(nn.Module):
    ''' Compose with three layers '''

    def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1):
        super(DecoderLayer, self).__init__()
        self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.enc_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout)
        self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)

    def forward(
            self, dec_input, enc_output,
            slf_attn_mask=None, dec_enc_attn_mask=None):
        dec_output, dec_slf_attn = self.slf_attn(
            dec_input, dec_input, dec_input, mask=slf_attn_mask)
        dec_output, dec_enc_attn = self.enc_attn(
            dec_output, enc_output, enc_output, mask=dec_enc_attn_mask)
        dec_output = self.pos_ffn(dec_output)
        return dec_output, dec_slf_attn, dec_enc_attn

基于 sort 算法的多目标跟踪

Tue, 09 Mar 2021 00:00:00 +0000

1. 概述

本文介绍基于 Sort 算法的多目标跟踪方案，实现车流量统计。该方案主要由三个部分组成。 yolov3 进行目标检测、使用匈牙利算法对目标进行关联、使用卡尔曼滤波器对跟踪目标进行修正。

2. yolov3 模型的使用

之前的文章中介绍过 yolov3 模型，这里不在进行赘述，本文使用 opencv 内部提供的 yolov3 接口进行实现，这里介绍下 yolov3 的 opencv 接口。

在 opencv 的 dnn 模块中包含了主流的深度学习模型，但要注意只提供推理功能，不支持模型训练。

2.1. yolov3 模型的应用

加载 yolov3 模型：

readNetFromDarknet: 参数 1 指定网络的配置信息，参数 2 指定模型的预训练权重。

weightsPath = "./yolo-coco/yoloV3.weights"
configPath = "./yolo-coco/yoloV3.cfg"
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)

获取网络输出层。

getLayerNames: 获取所有网络层的名称。
getUnconnectedOutLayers: 获取网络输出层的索引。

ln: 保存 yolov3 网络输出层，分别是 [yolo-82,yolo-94,yolo-106]

ln = net.getLayerNames()
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]

使用 opencv 接口获取视频流，并送入 yolov3 网络进行目标检测。

blobFromImage: dnn 接口，对图像进行加载并进行预处理。
net.forward: 将图像输入进网络，并进行前向传播，返回 yolov3 3 个层的目标预测结果。
遍历每一个输出层，遍历每一个预测结果，筛选 confidence > 0.3 的预测结果
boxes: 保存预测边界框的中心点坐标和宽高
confidences: 保存置信度
classIDs: 保存类别 id
NMSBoxes: 对上面进行筛选过后的预测信息通过非极大值抑制算法再进行筛选，其中 0.5 是 score_threshold ； 0.3 是 nms_threshold 。该函数返回符合要求的边界框的索引。

遍历经过非极大值抑制处理后的边界框，因为我们做车流量统计，只保留检测类别是车的边界框，将符合要求的边界框的左上角坐标和右下角坐标和置信度保存到 dets 中

vs = cv2.VideoCapture('./input/test_1.mp4')
(W, H) = (None, None)
writer = None
while True:
  (grabed, frame) = vs.read()
  if W is None or H is None:
      (H,W) = frame.shape[:2]
  blob = cv2.dnn.blobFromImage(frame, 1 / 255.0, (416, 416), swapRB=True, crop=False)
  net.setInput(blob)
  layerOutputs = net.forward(ln)
  boxes = []
  confidences = []
  classIDs = []
  for output in layerOutputs:
      for detection in output:
          # detction: [5:]: 表示类别，[0:4]: bbox 的位置信息， [4]: 置信度
          scores = detection[5:]
          classID = np.argmax(scores)
          confidence = scores[classID]

          if confidence > 0.3:
              # 将检测结果与原图片进行适配
              box = detection[0:4] * np.array([W, H, W, H])
              (centerX, centerY, width, height) = box.astype("int")
              # 左上角坐标
              x = int(centerX - width / 2)
              y = int(centerY - height / 2)
              # 更新目标框，置信度，类别
              boxes.append([x, y, int(width), int(height)])
              confidences.append(float(confidence))
              classIDs.append(classID)
  idxs = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.3)
  # 检测框:左上角和右下角
  dets = []
  if len(idxs) > 0:
      for i in idxs.flatten():
          if LABELS[classIDs[i]] == "car":
              (x, y) = (boxes[i][0], boxes[i][1])
              (w, h) = (boxes[i][2], boxes[i][3])
              # cv2.rectangle(frame,(x,y),(x+w,y+h),(0,255,0),2)
              dets.append([x, y, x + w, y + h, confidences[i]])

3. 匈牙利算法

匈牙利算法 (Hungarian Algorithm) 与 KM 算法 (Kuhn-Munkres Algorithm) 是用来解决多目标跟踪中的数据关联问题，匈牙利算法与 KM 算法都是为了求解二分图的最大匹配问题。

那什么是二分图呢？就是能分成两组 U 和 V ，其中， U 上的点不能相互连通，只能连接 V 中的点，同理， V 中的点不能相互连通，只能连去 U 中的点，这就是做二分图。

可以把二分图理解为视频中连续两帧中的所有检测框，第一帧所有检测框的集合称为 U ，第二帧所有检测框的集合称为 V 。同一帧的不同检测框不会为同一个目标，所以不需要互相关联，相邻两帧的检测框需要相互联通，最终将相邻两帧的检测框尽量完美地两两匹配起来。而求解这个问题的最优解就要用到匈牙利算法或者 KM 算法。

匈牙利算法和 KM 算法的原理这里就不在赘述了，相关的资料网上多的是，这里我们只需要知道他们是干什么的就可以了。其中 KM 算法是匈牙利算法的改进版本，它解决的是带权二分图的最优匹配问题。在多目标跟踪中目标关联是根据前一帧数据与后一帧数据的 IOU 作为权值进行关联的。

在 scipy 包中，通过 linear_sum_assignment 函数实现 KM 算法。

构造代价矩阵 cost
注意：传入 linear_sum_assignment 函数的代价矩阵要取负数。因为该方法的目的是代价最小，这里是求最大匹配，所以将 cost 取负数。 ``` from scipy.optimize import linear_sum_assignment import numpy as np

cost = np.array([[0.9,0.6,0,0],[0,0.3,0.9,0],[0.5,0.9,0,0],[0,0,0.2,0]]) row_ind,col_ind = linear_sum_assignment(-cost)

# 4.卡尔曼滤波
&#160; &#160; &#160; &#160;卡尔曼滤波无论是在单目标还是多目标领域都是很常用的一种算法，我们将卡尔曼滤波看做一种运动模型，用来对目标的位置进行预测，并且利用预测结果对跟踪的目标进行修正，属于自动控制理论中的一种方法。

&#160; &#160; &#160; &#160;在对视频中的目标进行跟踪时，当目标运动速度较慢时，很容易将前后两帧的目标进行关联，如下图所示:

![](/img/Kalman1.png)

&#160; &#160; &#160; &#160;如果目标运动速度比较快，或者进行隔帧检测时，在后续帧中，目标A已运动到前一帧B所在的位置，这时再进行关联就会得到错误的结果，将 A' 与 B 关联在一起。

&#160; &#160; &#160; &#160;那怎么才能避免这种出现关联误差呢？我们可以在进行目标关联之前，对目标在后续帧中出现的位置进行预测，然后与预测结果进行对比关联，如下图所示：

![](/img/Kalman2.png)

&#160; &#160; &#160; &#160;我们在对比关联之前，先预测出 A 和 B 在下一帧中的位置，然后再使用实际的检测位置与预测的位置进行对比关联，只要预测足够精确，几乎不会出现由于速度太快而关联错误的情况。

&#160; &#160; &#160; &#160;卡尔曼滤波就是用来预测目标在后续帧中出现的位置。卡尔曼滤波器最大的优点是采用递归的方法来解决线性滤波的问题，它只需要当前的测量值和前一个周期的预测值就能够进行状态估计。由于这种递归方法不需要大量的存储空间，每一步的计算量小，计算步骤清晰，非常适合计算机处理，因此卡尔曼滤波受到了普遍的欢迎，在各种领域具有广泛的应用前景。

&#160; &#160; &#160; &#160;简单理解卡尔曼滤波分为两步，预测和更新。预测是根据上一周期的预测值对当前状进行估计。更新是用当前的观测值更新卡尔曼滤波器，用于下次状态的估计，这实际上就是递归。更新阶段是卡尔曼滤波的数据融合，它融合了估计值和观测值的结果，充分利用两者的不确定性来得到更加准确的估计。

&#160; &#160; &#160; &#160;以上就是卡尔曼滤波的简单思想，如果想了解详细的推理过程可以自行搜集资料进行研究。

&#160; &#160; &#160; &#160;在实际使用卡尔曼滤波的时候，计算的步骤一般为：
* 预测阶段

![](/img/Kalman3.png)

* 更新阶段

![](/img/Kalman4.png)

# 4.1. 卡尔曼滤波器的实现
&#160; &#160; &#160; &#160;卡尔曼滤波器的实现在 filterpy 包里。 filterpy 是一个实现了各种滤波器的 Python 模块，它实现著名的卡尔曼滤波和粒子滤波器。我们可以直接调用该库完成卡尔曼滤波器实现。

from filterpy.kalman import KalmanFilter

&#160; &#160; &#160; &#160;定义一个卡尔曼滤波器用于跟踪目标边界框，上面已经提到了，卡尔曼的两大功能是预测和更新，该类中主要实现这两个接口。首先我们需要对滤波器进行初始化。

* 初始化卡尔曼滤波器的状态变量和观测输入，这里我们假设车的运动是一个等速模型。
* 状态变量 x 设定为一个七维向量， x = [u,v,s,r,u',v',s'].T ，分别表示目标中心位置的 x,y 坐标，面积 s 和当前目标框的纵横比，最后三个则是横向，纵向，面积的变化速率，其中速度部分初始化为 0 ，其他根据观测进行输入。
* 量测矩阵 H 是 4*7 的矩阵，将观测值与状态变量相对应。
* 根据经验值进行相应的协方差参数的设定 P Q R

class KalmanBoxTracker(object): count = 0

def __init__(self, bbox):
    # 定义等速模型
    # 内部使用KalmanFilter，7个状态变量和4个观测输入
    self.kf = KalmanFilter(dim_x=7, dim_z=4)
    # F是状态变换模型
    self.kf.F = np.array(
        [[1, 0, 0, 0, 1, 0, 0], [0, 1, 0, 0, 0, 1, 0], [0, 0, 1, 0, 0, 0, 1], [0, 0, 0, 1, 0, 0, 0],
         [0, 0, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 0, 1]])
    # H是观测函数
    self.kf.H = np.array(
        [[1, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0], [0, 0, 0, 1, 0, 0, 0]])
    # R是观测函数
    self.kf.R[2:, 2:] *= 10.
    # P是协方差矩阵
    self.kf.P[4:, 4:] *= 1000.  # give high uncertainty to the unobservable initial velocities
    self.kf.P *= 10.
    # Q是过程噪声矩阵
    self.kf.Q[-1, -1] *= 0.01
    self.kf.Q[4:, 4:] *= 0.01
    # 内部状态估计
    self.kf.x[:4] = convert_bbox_to_z(bbox)
    self.time_since_update = 0
    self.id = KalmanBoxTracker.count
    KalmanBoxTracker.count += 1
    self.history = []
    self.hits = 0
    self.hit_streak = 0
    self.age = 0 ```

使用观测值更新状态变量。

time_since_update: 该变量用于记录当前卡尔曼滤波器跟踪的目标有多少次未被跟踪到，当超过指定次数后，将该目标从跟踪列表中删除。每次 update 被调用时将该变量清零。

hit_streak: 该变量用于记录当前卡尔曼滤波器跟踪的目标被成功跟踪了多少次，当超过一定次数后，认为该目标跟踪成功

def update(self, bbox):
  self.time_since_update = 0
  self.history = []
  self.hits += 1
  self.hit_streak += 1
  self.kf.update(convert_bbox_to_z(bbox))

使用前一状态对当前状太进行估计。

time_since_update: 该变量用于记录当前卡尔曼滤波器跟踪的目标有多少次未被跟踪到，当超过指定次数后，将该目标从跟踪列表中删除。每次 predict 被调用时将该变量加 1 。

hit_streak: 该变量用于记录当前卡尔曼滤波器跟踪的目标被成功跟踪了多少次，当超过一定次数后，认为该目标跟踪成功，当有一次为跟踪到时，将该变量置 0 。

def predict(self):
  if (self.kf.x[6] + self.kf.x[2]) <= 0:
      self.kf.x[6] *= 0.0
  self.kf.predict()
  self.age += 1
  if self.time_since_update > 0:
      self.hit_streak = 0
  self.time_since_update += 1
  self.history.append(convert_x_to_bbox(self.kf.x))
  return self.history[-1]

返回当前估计值。

def get_state(self):
    return convert_x_to_bbox(self.kf.x)

5. Sort 算法实现

Sort 算法实际上是一个多目标跟踪器，管理多个 KalmanBoxTracker 对象。

首先我们看 Sort 类的构造方法。

max_age: 目标未被检测到的帧数，超过之后会被删除

min_hits: 目标连续被检测到 min_hits 次才对目标进行跟踪

class Sort(object):
  def __init__(self, max_age=1, min_hits=3):
      self.max_age = max_age
      self.min_hits = min_hits
      self.trackers = []   # ？
      self.frame_count = 0  # ？

update 方法是 Sort 算法的核心，它实现了对目标的跟踪。

dets: 输入 yolov3 网络检测到的目标边界框
self.frame_count: 记录视频帧数
trks: 存储跟踪器的预测，根据当前所有的卡尔曼跟踪器个数创建二维数组，行号为卡尔曼滤波器的标识索引，列向量为跟踪框的位置和 ID ，当第一次开始检测时 self.trackers 的长度是 0
to_del: 存储要删除的目标框
ret: 存储要返回的追踪目标框
循环遍历卡尔曼跟踪器列表，如果是第一次开始检测不会进入 for 循环。首先使用卡尔曼滤波器预测目标当前的位置，并把预测值保存到 trks 数组中，如果跟踪框中包含空值则将该跟踪框添加到要删除的列表中。
将 trks 中存在无效值的行删除
逆向删除异常的跟踪器，防止破坏索引
将目标检测框与卡尔曼滤波器预测的跟踪框关联获取跟踪成功的目标，新增的目标，离开画面的目标。
用跟踪成功的目标框更新到对应的卡尔曼滤波器，实际上就是之前提到的用观测值更新卡尔曼滤波器。
为新增的目标创建新的卡尔曼滤波器对象进行跟踪。并将卡尔曼滤波器添加到 self.trackers 中。
反向遍历卡尔曼滤波器列表，获取每一个卡尔曼滤波器的估计值，判断跟踪目标是否成功，成功则添加到列表中。判断是否跟踪成功的条件是连续 self.min_hits 次都跟踪到该目标。如果超过 self.max_age 次未跟踪到目标则认为该目标离开画面，删除对应的卡尔曼滤波器。

最终将跟踪结果拼接到一起返回。

def update(self, dets):
  self.frame_count += 1
  trks = np.zeros((len(self.trackers), 5))
  to_del = []
  ret = []
  for t, trk in enumerate(trks):
      pos = self.trackers[t].predict()[0]
      trk[:] = [pos[0], pos[1], pos[2], pos[3], 0]
      if np.any(np.isnan(pos)):
          to_del.append(t)
  trks = np.ma.compress_rows(np.ma.masked_invalid(trks))
  for t in reversed(to_del):
      self.trackers.pop(t)
  matched, unmatched_dets, unmatched_trks = associate_detections_to_trackers(dets, trks)

  for t, trk in enumerate(self.trackers):
      if t not in unmatched_trks:
          d = matched[np.where(matched[:, 1] == t)[0], 0]
          trk.update(dets[d, :][0])

  for i in unmatched_dets:
      trk = KalmanBoxTracker(dets[i, :])
      self.trackers.append(trk)

  i = len(self.trackers)
  for trk in reversed(self.trackers):
      d = trk.get_state()[0]
      if (trk.time_since_update < 1) and (trk.hit_streak >= self.min_hits or self.frame_count <= self.min_hits):
          ret.append(np.concatenate((d, [trk.id + 1])).reshape(1, -1))
      i -= 1
      if trk.time_since_update > self.max_age:
          self.trackers.pop(i)
  if len(ret) > 0:
      return np.concatenate(ret)
  return np.empty((0, 5))

在目标跟踪中，需要使用 KM 算法将卡尔曼滤波器估计的边界框与 yolov3 模型检测的边界框进行关联。

detections: 表示 yolov3 模型检测的边界框
trackers: 卡尔曼滤波器跟踪的边界框
iou_threshold: iou 阈值，该阈值用来判断估计模型与检测模型的匹配程度。
函数返回跟踪成功目标的矩阵: matchs ；新增目标的矩阵: unmatched_detections ；跟踪失败即离开画面的目标矩阵: unmatched_trackers
第一次进来时，跟踪目标为 0 ，直接返回。
计算检测框和估计框的 iou ，生成一个 [len(detections), len(trackers)] 大小的矩阵
linear_assignment: 使用 KM 算法计算匹配结果，将匹配结果保存到 matched_indices 中。
记录未匹配的检测框及跟踪框，未匹配的检测框放入 unmatched_detections 中，表示有新的目标进入画面，要新增跟踪器跟踪目标；未匹配的跟踪框放入 unmatched_trackers 中，表示目标离开之前的画面，应删除对应的跟踪器

将匹配成功的跟踪框放入 matches 中，要求 iou 大与设定的阈值。将低于该阈值的分别放入 unmatched_detections 和 unmatched_trackers 中。

def associate_detections_to_trackers(detections, trackers, iou_threshold=0.3):
  if (len(trackers) == 0) or (len(detections) == 0):
      return np.empty((0, 2), dtype=int), np.arange(len(detections)), np.empty((0, 5), dtype=int)

  iou_matrix = np.zeros((len(detections), len(trackers)), dtype=np.float32)
  for d, det in enumerate(detections):
      for t, trk in enumerate(trackers):
          iou_matrix[d, t] = iou(det, trk)

  result = linear_sum_assignment(-iou_matrix)
  matched_indices = np.array(list(zip(*result)))

  unmatched_detections = []
  for d, det in enumerate(detections):
      if d not in matched_indices[:, 0]:
          unmatched_detections.append(d)

  unmatched_trackers = []
  for t, trk in enumerate(trackers):
      if t not in matched_indices[:, 1]:
          unmatched_trackers.append(t)

  matches = []
  for m in matched_indices:
      if iou_matrix[m[0], m[1]] < iou_threshold:
          unmatched_detections.append(m[0])
          unmatched_trackers.append(m[1])
      else:
          matches.append(m.reshape(1, 2))
  if len(matches) == 0:
      matches = np.empty((0, 2), dtype=int)
  else:
      matches = np.concatenate(matches, axis=0)

  return matches, np.array(unmatched_detections), np.array(unmatched_trackers)

如何计算 mAP

Sun, 10 Jan 2021 00:00:00 +0000

1. 概述

mAP 是评价目标检测模型模型好坏的重要指标，这篇文章就介绍下如何计算 mAP 。

2. 基本概念

介绍 mAP 计算之前首先我们先了解几个重要概念。

2.1. IOU

IOU 是 Intersection over Union 的缩写，意为交并比。它用来衡量真实边界框与预测边界框的重合程度，它的计算公式为 [交集/并集] 。

2.2. 精确率和召回率

介绍精确率和召回率的之前先了解 TP 、 TN 、 FP 、 FN 这几个概念。

TP: True Positive 的缩写，意思是预测为正样本，实际上也为正样本。

TN: True Negtive 的缩写，意思是预测为负样本，实际上也是负样本。

FP: False Positive 的缩写，意思是预测为正样本，实际上是负样本。

FN: False Negtive 的缩写，意思是预测为负样本，实际上是正样本。

以上这四个值很好理解，就是非常容易搞混，可以这样记： T 和 F 代表的是该样本是否被正确分类； P 和 N 代表的是该样本被预测成了正样本还是负样本。

有了以上概念之后，精确率和召回率就很好理解了。

精确率 = TP / TP + FP : 表示分类器认为是正类并且确实是正类的部分占所有分类器认为是正类的比例。

召回率 = TP / TP + FN : 表示分类器认为是正类并且确实是正类的部分占所有确实是正类的比例。

2.3. 置信度

置信度为 pr(object)*iou(b,object) ，表示预测边界框是否包含物体与物体与真实边界框的 IOU 的乘积。置信度用来表示模型输出的可信度，如果置信度设置的高的话，预测的结果和实际情况就很符合，如果置信度低的话，就会有很多误检测。

3. AP 的引入

为什么要引入 AP 这个概念呢？让我们看下对于一个模型的好坏如果只使用精确率或者召回率会有什么问题。假设一幅图像里面总共有 3 个正样本，目标检测对这幅图的预测结果有 10 个，其中 3 个实际上是正样本， 7 个实际上是负样本。对应置信度如图所示：

如果我们将可以接受的置信度设置为 0.95 的话，那么目标检测算法就会将序号为 1 的样本作为正样本，其它的都是负样本。此时 TP = 1 ， FP = 0 ， FN = 2 。那么精确率 = 1 ，召回率 = 1/3 。

此时精确率非常高，但是事实上我们只检测出一个正样本，还有两个没有检测出来，因此只用精确率就不能很好的表示模型的好坏。

如果我们将可以接受的置信度设置为 0.35 时，目标检测算法就会将序号为 1 到 6 的样本都作为正样本，其它的是负样本。此时 TP = 3 ， FP = 3 ， FN = 0 。那么精确率 = 1/2 ，召回率 = 1 。

此时召回率非常高，但是事实上目标检测算法认为是正样本的样本里面，有 3 个样本确实是正样本，但有 3 个是负样本，存在非常严重的误检测，因此只用召回率也不恶嗯呐很好的表示模式的好坏。

基于以上单个指标的局限性，引入了 AP ， AP 指的是利用不同的精确率和召回率的点的组合，画出来的曲线下面的面积。当我们取不同的置信度，可以获得不同的精确率和召回率，当我们取得置信度够密集的时候，就可以获得非常多的精确率和召回率。此时精确率和召回率可以在图片上画出一条线，这条线下部分的面积就是某个类的 AP 值。 mAP 就是所有的类的 AP 值求平均。

4. AP 的代码实现

这里介绍的参考代码来自 Github 。在使用这个代码进行 mAP 计算之前需要做些准备工作。

首先，准备预测结果，并放到 detection-results 中。

然后，准备标签文件，并放到 ground-truth 中。

最后，准备图片文件，并放到 image-optional 中。这个目录用来做可视化，可以没有。

这里我先假设以上文件已准备好，并且这里假设不做可视化，下面我们直接看代码实现：

首先初始化已经准备好了的目录。

os.chdir(os.path.dirname(os.path.abspath(__file__)))
GT_PATH = os.path.join(os.getcwd(), 'input', 'ground-truth')
DR_PATH = os.path.join(os.getcwd(), 'input', 'detection-results')
IMG_PATH = os.path.join(os.getcwd(), 'input', 'images-optional')

创建两个目录， .temp_files 是临时目录，存放计算过程中的临时数据； output 存放最终的计算结果。

TEMP_FILES_PATH = ".temp_files"
if not os.path.exists(TEMP_FILES_PATH):
  os.makedirs(TEMP_FILES_PATH)
output_files_path = "output"
if os.path.exists(output_files_path):
  shutil.rmtree(output_files_path)
os.makedirs(output_files_path)

接下来是遍历所有标签文件，将每个标签文件里的目标解析出来，放到 _ground_truth.json 中。

获得所有标签文件的列表，然后排序。
gt_counter_per_class: 记录每个类别标签的个数。

counter_images_per_class: 记录每一类别存在于多少张图像中。

ground_truth_files_list = glob.glob(GT_PATH + '/*.txt')
ground_truth_files_list.sort()
gt_counter_per_class = {}
counter_images_per_class = {}

file_id: 去掉后缀的文件名
temp_path: file_id 对应的预测值文件
lines_list: 读取标签文件的每一行
bounding_boxes: 存储标签信息
is_difficult: 目标是否是难检测样本
already_seen_classes: 某一类别的样本是否在一张图像中多次出现
解析标签文件的每一行，格式为 class_name, left, top, right, bottom, _difficult ， _difficult 可能没有
将解析出来的信息添加到 bounding_boxes 中
难检测样本不记录到个数里
gt_counter_per_class: 记录每一个类总共有多少个样本
already_seen_classes: 某一类别的样本是否在一张图像中多次出现，如果某一类别在一张图像中没有出现过，则 counter_images_per_class 加 1 ， counter_images_per_class: 记录每一类别存在于多少张图像中

将解析出来的信息写入 file_id + _ground_truth.json 文件中

gt_files = []
for txt_file in ground_truth_files_list:
  file_id = txt_file.split(".txt", 1)[0]
  file_id = os.path.basename(os.path.normpath(file_id))
  # check if there is a correspondent detection-results file
  temp_path = os.path.join(DR_PATH, (file_id + ".txt"))
  lines_list = file_lines_to_list(txt_file)
  bounding_boxes = []
  is_difficult = False
  already_seen_classes = []
  for line in lines_list:
      if "difficult" in line:
              class_name, left, top, right, bottom, _difficult = line.split()
              is_difficult = True
      else:
              class_name, left, top, right, bottom = line.split()
      bbox = left + " " + top + " " + right + " " +bottom
      if is_difficult:
          bounding_boxes.append({"class_name":class_name, "bbox":bbox, "used":False, "difficult":True})
          is_difficult = False
      else:
          bounding_boxes.append({"class_name":class_name, "bbox":bbox, "used":False})
          if class_name in gt_counter_per_class:
              gt_counter_per_class[class_name] += 1
          else:
              gt_counter_per_class[class_name] = 1

          if class_name not in already_seen_classes:
              if class_name in counter_images_per_class:
                  counter_images_per_class[class_name] += 1
              else:
                  counter_images_per_class[class_name] = 1
              already_seen_classes.append(class_name)


  new_temp_file = TEMP_FILES_PATH + "/" + file_id + "_ground_truth.json"
  gt_files.append(new_temp_file)
  with open(new_temp_file, 'w') as outfile:
      json.dump(bounding_boxes, outfile)

记录所有的类别并排序，保存类别个数

gt_classes = list(gt_counter_per_class.keys())
gt_classes = sorted(gt_classes)
n_classes = len(gt_classes)

遍历所有预测信息文件，将预测为同一类别的信息保存到一个文件中。

获得预测信息文件，并排序

dr_files_list = glob.glob(DR_PATH + '/*.txt')
dr_files_list.sort()

遍历每一个类别，统计所有预测类别为该类别的信息，将信息保存到 bounding_boxes 中
将 bounding_boxes 按照置信度排序

将每个类别的所有信息保存到 class_name + _dr.json 文件中

for class_index, class_name in enumerate(gt_classes):
  bounding_boxes = []
  for txt_file in dr_files_list:
      file_id = txt_file.split(".txt",1)[0]
      file_id = os.path.basename(os.path.normpath(file_id))
      temp_path = os.path.join(GT_PATH, (file_id + ".txt"))
      lines = file_lines_to_list(txt_file)
      for line in lines:
          tmp_class_name, confidence, left, top, right, bottom = line.split()
          if tmp_class_name == class_name:
              bbox = left + " " + top + " " + right + " " +bottom
              bounding_boxes.append({"confidence":confidence, "file_id":file_id, "bbox":bbox})
  bounding_boxes.sort(key=lambda x:float(x['confidence']), reverse=True)
  with open(TEMP_FILES_PATH + "/" + class_name + "_dr.json", 'w') as outfile:
      json.dump(bounding_boxes, outfile)

准备工作做完了，开始计算 mAP 。

将 AP 值写入 output.txt 文件
count_true_positives: 遍历每个类别，统计每个类别的 tp 个数。
取出之前保存的每个类别的预测信息
遍历每一条预测信息
根据 file_id 找到对应的标签信息
获取预测的边界框信息和标签的边界框信息，计算 IOU 值，找到最佳的 iou 值
如果 iou 大于等于给定的最小 iou 阈值，且标签信息不是难预测的样本，表示该预测为正类且预测正确
gt_match[“used”]=False: 表示该标签被用过后就不能在重复使用，如果已被使用，则为假正类 fp[idx] = 1
如果 iou 小于给定的最小 iou 阈值，则为假正类 fp[idx] = 1
计算精确率和召回率，统计假正类的个数，统计真正类的个数
召回率=真正类个数/所有正样本的个数
精确率=真正类个数/所有预测为正类的个数
AP 为召回率和精确率围成的曲线的面积，使用 voc_ap 函数计算各个类的 AP

mAP = sum_AP / n_classes: mAP 是所有类别 AP 的平均值。

sum_AP = 0.0
ap_dictionary = {}
lamr_dictionary = {}
with open(output_files_path + "/output.txt", 'w') as output_file:
  output_file.write("# AP and precision/recall per class\n")
  count_true_positives = {}
  # 遍历每个类别
  for class_index, class_name in enumerate(gt_classes):
      count_true_positives[class_name] = 0

      dr_file = TEMP_FILES_PATH + "/" + class_name + "_dr.json"
      dr_data = json.load(open(dr_file))


      nd = len(dr_data)
      tp = [0] * nd # creates an array of zeros of size nd
      fp = [0] * nd
      for idx, detection in enumerate(dr_data):
          file_id = detection["file_id"]

          gt_file = TEMP_FILES_PATH + "/" + file_id + "_ground_truth.json"
          ground_truth_data = json.load(open(gt_file))
          ovmax = -1
          gt_match = -1

          bb = [ float(x) for x in detection["bbox"].split() ]
          for obj in ground_truth_data:

              if obj["class_name"] == class_name:
                  bbgt = [ float(x) for x in obj["bbox"].split() ]
                  bi = [max(bb[0],bbgt[0]), max(bb[1],bbgt[1]), min(bb[2],bbgt[2]), min(bb[3],bbgt[3])]
                  iw = bi[2] - bi[0] + 1
                  ih = bi[3] - bi[1] + 1
                  if iw > 0 and ih > 0:
                      # compute overlap (IoU) = area of intersection / area of union
                      ua = (bb[2] - bb[0] + 1) * (bb[3] - bb[1] + 1) + (bbgt[2] - bbgt[0]
                                      + 1) * (bbgt[3] - bbgt[1] + 1) - iw * ih
                      ov = iw * ih / ua
                      if ov > ovmax:
                          ovmax = ov
                          gt_match = obj

          # set minimum overlap
          min_overlap = MINOVERLAP

          if ovmax >= min_overlap:
              if "difficult" not in gt_match:
                      if not bool(gt_match["used"]):
                          # true positive
                          tp[idx] = 1
                          gt_match["used"] = True
                          count_true_positives[class_name] += 1
                          # update the ".json" file
                          with open(gt_file, 'w') as f:
                                  f.write(json.dumps(ground_truth_data))
                      else:
                          # false positive (multiple detection)
                          fp[idx] = 1
          else:
              # false positive
              fp[idx] = 1
              if ovmax > 0:
                  status = "INSUFFICIENT OVERLAP"

      # compute precision/recall
      cumsum = 0
      for idx, val in enumerate(fp):
          fp[idx] += cumsum
          cumsum += val
      cumsum = 0
      for idx, val in enumerate(tp):
          tp[idx] += cumsum
          cumsum += val
      #print(tp)
      rec = tp[:]
      # 
      for idx, val in enumerate(tp):
          rec[idx] = float(tp[idx]) / gt_counter_per_class[class_name]
      #print(rec)
      prec = tp[:]
      # 
      for idx, val in enumerate(tp):
          prec[idx] = float(tp[idx]) / (fp[idx] + tp[idx])
      #print(prec)

      ap, mrec, mprec = voc_ap(rec[:], prec[:])
      sum_AP += ap
      text = "{0:.2f}%".format(ap*100) + " = " + class_name + " AP " #class_name + " AP = {0:.2f}%".format(ap*100)
      """
       Write to output.txt
      """
      rounded_prec = [ '%.2f' % elem for elem in prec ]
      rounded_rec = [ '%.2f' % elem for elem in rec ]
      output_file.write(text + "\n Precision: " + str(rounded_prec) + "\n Recall :" + str(rounded_rec) + "\n\n")
      if not args.quiet:
          print(text)
      ap_dictionary[class_name] = ap

      n_images = counter_images_per_class[class_name]
      lamr, mr, fppi = log_average_miss_rate(np.array(prec), np.array(rec), n_images)
      lamr_dictionary[class_name] = lamr

  output_file.write("\n# mAP of all classes\n")
  mAP = sum_AP / n_classes
  text = "mAP = {0:.2f}%".format(mAP*100)
  output_file.write(text + "\n")
  print(text)

深度学习之(十四)EfficientNet 网络

Sat, 10 Oct 2020 00:00:00 +0000

1. 概述

EfficientNet 是谷歌在 2019 年提出的新的特征提取网络。它的主要创新点并不是结构，不像 ResNet 、 SENet 发明了 shortcut 或 attention 机制， EfficientNet 的 base 结构是利用结构搜索搜出来的，然后使用 compound scaling 规则放缩，得到一系列表现优异的网络： B0~B7 。

增加网络参数可以获得更好的精度（有足够的数据，不过拟合的条件下），例如 ResNet 可以加深从 ResNet-18 到 ResNet-200 。增加网络参数的方式有三种：深度、宽度和分辨率。深度是指网络的层数，宽度指网络中卷积的 channel 数，分辨率是指通过网络输入大小（例如从 112x112 到 224x224 ）。直观上来讲，这三种缩放方式并不独立。对于分辨率高的图像，应该用更深的网络，因为需要更大的感受野，同时也应该增加网络宽度来获得更细粒度的特征。之前增加网络参数都是单独放大这三种方式中的一种，并没有同时调整，也没有调整方式的研究。 EfficientNet 使用了 compound scaling 方法，统一缩放网络深度、宽度和分辨率。

2. 网络结构

EfficientNet B0~B7 网络结构由三个部分组成，分别是 Stem 、 MBConvBlock 和 Final Layers 。其中 Stem 就是标准的卷积、 BN 、激活函数。 MBConvBlock 结构类似 MobileNetV3 的网络结构。 Final Layers 只是在 1x1 的卷积加上预测器。

3. 网络实现

首先，我们先看一下网络的超参数。

params 前三个参数定义了 EfficientNet B0~B7 网络结构在三个维度的缩放比例，最后一个参数是 dropout_rate 。

params = {
  'efficientnet_b0': (1.0, 1.0, 224, 0.2),
  'efficientnet_b1': (1.0, 1.1, 240, 0.2),
  'efficientnet_b2': (1.1, 1.2, 260, 0.3),
  'efficientnet_b3': (1.2, 1.4, 300, 0.3),
  'efficientnet_b4': (1.4, 1.8, 380, 0.4),
  'efficientnet_b5': (1.6, 2.2, 456, 0.4),
  'efficientnet_b6': (1.8, 2.6, 528, 0.5),
  'efficientnet_b7': (2.0, 3.1, 600, 0.5),
}

settings 变量定一个 MBConv 网络的结构，其中 t 表示输入通道的扩张系数；
c 表示输出通道数；
n 表示该模块的重复次数；
s 表示 stride ；

k 表示卷积核大小；

settings = [
      # t,  c, n, s, k
      [1,  16, 1, 1, 3],  # MBConv1_3x3, SE, 112 -> 112
      [6,  24, 2, 2, 3],  # MBConv6_3x3, SE, 112 ->  56
      [6,  40, 2, 2, 5],  # MBConv6_5x5, SE,  56 ->  28
      [6,  80, 3, 2, 3],  # MBConv6_3x3, SE,  28 ->  14
      [6, 112, 3, 1, 5],  # MBConv6_5x5, SE,  14 ->  14
      [6, 192, 4, 2, 5],  # MBConv6_5x5, SE,  14 ->   7
      [6, 320, 1, 1, 3]   # MBConv6_3x3, SE,   7 ->   7
  ]

3.1. Stem 结构实现

Stem 就是标准卷积、 BN 、激活函数结构。激活函数用的是 swish 。

class ConvBNReLU(nn.Sequential):

    def __init__(self, in_planes, out_planes, kernel_size, stride=1, groups=1):
        padding = self._get_padding(kernel_size, stride)
        super(ConvBNReLU, self).__init__(
            nn.ZeroPad2d(padding),
            nn.Conv2d(in_planes, out_planes, kernel_size, stride, padding=0, groups=groups, bias=False),
            nn.BatchNorm2d(out_planes),
            Swish(),
        )

    def _get_padding(self, kernel_size, stride):
        p = max(kernel_size - stride, 0)
        return [p // 2, p - p // 2, p // 2, p - p // 2]

3.2. MBConvBlock 结构实现

MBConvBlock 结构实现如下：

self.use_residual: 在输入和输出相等，且 stride==1 的时候，增加残差网络
使用 1x1 卷积升维
进行深度可分离卷积
增加注意力机制
使用 1x1 卷积降维到输出维度
最后增加 BN 层

前向传播时，如果使用 use_residual ，则要使用 _drop_connect

class MBConvBlock(nn.Module):

  def __init__(self,
               in_planes,
               out_planes,
               expand_ratio,
               kernel_size,
               stride,
               reduction_ratio=4,
               drop_connect_rate=0.2):
      super(MBConvBlock, self).__init__()
      self.drop_connect_rate = drop_connect_rate
      self.use_residual = in_planes == out_planes and stride == 1
      assert stride in [1, 2]
      assert kernel_size in [3, 5]

      hidden_dim = in_planes * expand_ratio
      reduced_dim = max(1, int(in_planes / reduction_ratio))

      layers = []
      # pw
      if in_planes != hidden_dim:
          layers += [ConvBNReLU(in_planes, hidden_dim, 1)]

      layers += [
          # dw
          ConvBNReLU(hidden_dim, hidden_dim, kernel_size, stride=stride, groups=hidden_dim),
          # se
          SqueezeExcitation(hidden_dim, reduced_dim),
          # pw-linear
          nn.Conv2d(hidden_dim, out_planes, 1, bias=False),
          nn.BatchNorm2d(out_planes),
      ]

      self.conv = nn.Sequential(*layers)

  def _drop_connect(self, x):
      if not self.training:
          return x
      keep_prob = 1.0 - self.drop_connect_rate
      batch_size = x.size(0)
      random_tensor = keep_prob
      random_tensor += torch.rand(batch_size, 1, 1, 1, device=x.device)
      binary_tensor = random_tensor.floor()
      return x.div(keep_prob) * binary_tensor

  def forward(self, x):
      if self.use_residual:
          return x + self._drop_connect(self.conv(x))
      else:
          return self.conv(x)

3.3. Final Layers 结构实现

Final Layers 结构是 1x1 卷积和全链接层输出为类别的个数。

3.4. EfficientNet 网络的整体实现

初始化一个 EfficientNet 网络，传入参数为网络名称，例如： efficientnet_b0 。

def efficientnet_b0(pretrained=False, progress=True, **kwargs):
    return _efficientnet('efficientnet_b0', pretrained, progress, **kwargs)

def _efficientnet(arch, pretrained, progress, **kwargs):
    width_mult, depth_mult, _, dropout_rate = params[arch]
    model = EfficientNet(width_mult, depth_mult, dropout_rate, **kwargs)
    if pretrained:
        state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)

        if 'num_classes' in kwargs and kwargs['num_classes'] != 1000:
            del state_dict['classifier.1.weight']
            del state_dict['classifier.1.bias']

        model.load_state_dict(state_dict, strict=False)
    return model

EfficientNet 网络中根据不同的网络名称对网络的宽度和深度进行缩放。

宽度缩放： _round_filters ，缩放后的数字要能被 8 整除
深度缩放： _round_repeats ，缩放后的数字向上取整

class EfficientNet(nn.Module):

def __init__(self, width_mult=1.0, depth_mult=1.0, dropout_rate=0.2, num_classes=1000):
    super(EfficientNet, self).__init__()

    # yapf: disable
    settings = [
        # t,  c, n, s, k
        [1,  16, 1, 1, 3],  # MBConv1_3x3, SE, 112 -> 112
        [6,  24, 2, 2, 3],  # MBConv6_3x3, SE, 112 ->  56
        [6,  40, 2, 2, 5],  # MBConv6_5x5, SE,  56 ->  28
        [6,  80, 3, 2, 3],  # MBConv6_3x3, SE,  28 ->  14
        [6, 112, 3, 1, 5],  # MBConv6_5x5, SE,  14 ->  14
        [6, 192, 4, 2, 5],  # MBConv6_5x5, SE,  14 ->   7
        [6, 320, 1, 1, 3]   # MBConv6_3x3, SE,   7 ->   7
    ]
    # yapf: enable

    out_channels = _round_filters(32, width_mult)
    features = [ConvBNReLU(3, out_channels, 3, stride=2)]

    in_channels = out_channels
    for t, c, n, s, k in settings:
        out_channels = _round_filters(c, width_mult)
        repeats = _round_repeats(n, depth_mult)
        for i in range(repeats):
            stride = s if i == 0 else 1
            features += [MBConvBlock(in_channels, out_channels, expand_ratio=t, stride=stride, kernel_size=k)]
            in_channels = out_channels

    last_channels = _round_filters(1280, width_mult)
    features += [ConvBNReLU(in_channels, last_channels, 1)]

    self.features = nn.Sequential(*features)
    self.classifier = nn.Sequential(
        nn.Dropout(dropout_rate),
        nn.Linear(last_channels, num_classes),
    )

    # weight initialization
    for m in self.modules():
        if isinstance(m, nn.Conv2d):
            nn.init.kaiming_normal_(m.weight, mode='fan_out')
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.BatchNorm2d):
            nn.init.ones_(m.weight)
            nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Linear):
            fan_out = m.weight.size(0)
            init_range = 1.0 / math.sqrt(fan_out)
            nn.init.uniform_(m.weight, -init_range, init_range)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

def forward(self, x):
    x = self.features(x)
    x = x.mean([2, 3])
    x = self.classifier(x)
    return x

深度学习之(十三)MobildNet 网络

Sat, 01 Aug 2020 00:00:00 +0000

1. 概述

MobileNet 是谷歌 2017 提出的用于移动设备上的轻量级神经网络。那么为什么 MobileNet 是如何做到在不影响模型精度的条件下，大幅减少模型参数的呢？答案是深度可分离卷积。在随后的两年谷歌又推出了 V2 和 V3 版本。

2. 深度可分离卷积

深度可分离卷积就是将普通卷积拆分成为一个深度卷积和一个逐点卷积。

2.1. 深度卷积

对于标准卷积，输入一个 12×12×3 的一个图像，经过一个卷积核大小为 5×5×3 的卷积得到一个 8×8×1 的输出特征图。如果有 256 个卷积核，我们将会得到一个 8×8×256 的输出特征图。

而对于深度卷积，是将卷积核拆分成为单通道形式，在不改变输入特征图像的深度的情况下，对每一通道进行卷积操作，这样就得到了和输入特征图通道数一致的输出特征图。输入 12×12×3 的特征图，经过 5×5×1×3 的深度卷积之后，得到了 8×8×3 的输出特征图。输入和输出的维度是不变的。

这样就会有一个问题，通道数太少，特征图的维度太少，能获取到足够的有效信息吗？这时逐点卷积就该登场了。

2.2. 逐点卷积

逐点卷积就是 1×1 卷积。主要作用就是对特征图进行升维和降维。

在深度卷积的过程中，我们得到了 8×8×3 的输出特征图，我们用 256 个 1×1×3 的卷积核对输入特征图进行卷积操作，输出的特征图和标准的卷积操作一样都是 8×8×256 了。

以上述为例，这里我们对比一下标准卷积和深度可分离卷积参数量的大小，先看一下标准卷积：

5×5×3×256 = 19200

对于深度可分离卷积参数量为：

5x5x1x3 + 1x1x3x256 = 843

很明显，深度可分离卷积的参数量相比普通卷积的参数量有了大幅下降。

3. MobileNet V1

3.1. 网络结构

MobileNet V1 网络主干结构如下图所示：

MobileNet V1 网络结构组成如下图所示，其中 s1 表示步长为 1 ， s2 表示步长为 2 ， Conv 表示普通卷积网络， Conv dw 表示深度可分离卷积网络中的深度卷积。

MobileNet V1 网络参数和计算量分布：

3.1. 网络实现

我们从下面代码可以看出 MobileNetV1 网络的结构就是普通卷积和深度可分离卷积的堆叠。
conv_bn 和 conv_dw 分别实现了普通卷积和深度可分离卷积，输入参数分别为输入通道数、输出通道数和步长。
深度可分离卷积中的深度卷积将 groups 设为 inp 。

注意： conv_dw 函数实现了网络结构里的 Conv dw 和 Conv 两个结构

class MobileNetV1(nn.Module):
  def __init__(self):
      super(Net, self).__init__()

      def conv_bn(inp, oup, stride):
          return nn.Sequential(
              nn.Conv2d(inp, oup, 3, stride, 1, bias=False),
              nn.BatchNorm2d(oup),
              nn.ReLU(inplace=True)
          )

      def conv_dw(inp, oup, stride):
          return nn.Sequential(
              nn.Conv2d(inp, inp, 3, stride, 1, groups=inp, bias=False),
              nn.BatchNorm2d(inp),
              nn.ReLU(inplace=True),
    
              nn.Conv2d(inp, oup, 1, 1, 0, bias=False),
              nn.BatchNorm2d(oup),
              nn.ReLU(inplace=True),
          )

      self.model = nn.Sequential(
          conv_bn(  3,  32, 2), 
          conv_dw( 32,  64, 1),
          conv_dw( 64, 128, 2),
          conv_dw(128, 128, 1),
          conv_dw(128, 256, 2),
          conv_dw(256, 256, 1),
          conv_dw(256, 512, 2),
          conv_dw(512, 512, 1),
          conv_dw(512, 512, 1),
          conv_dw(512, 512, 1),
          conv_dw(512, 512, 1),
          conv_dw(512, 512, 1),
          conv_dw(512, 1024, 2),
          conv_dw(1024, 1024, 1),
          nn.AvgPool2d(7),
      )
      self.fc = nn.Linear(1024, 1000)

  def forward(self, x):
      x = self.model(x)
      x = x.view(-1, 1024)
      x = self.fc(x)
return x

4. MobileNet V2

4.1. 网络结构

在 MobileNet V1 网络中，利用 3×3 的深度可分离卷积提取特征，然后利用 1×1 的卷积来扩张通道。这样网络结构既能减少不小的参数量、计算量，提高了网络运算速度，又能的得到一个接近于标准卷积的还不错的结果，看起来是很美好的。但是，实际使用的时候，发现深度卷积部分的卷积核比较容易训废掉，即训完之后发现深度卷积训出来的卷积核有不少是空的。作者认为这是 ReLU 激活函数导致的。

从上图看出，当 n = 2,3 时，与 Input 相比有很大一部分的信息已经丢失了。而当 n = 15 到 30 时，还是有相当多的地方被保留了下来。也就是说，对低维度做 ReLU 运算，很容易造成信息的丢失。而在高维度进行 ReLU 运算的话，信息的丢失则会很少。这就解释了为什么深度卷积的卷积核有不少是空。既然是 ReLU 导致的上述问题，那么就将 ReLU 替换成线性激活函数。于是就有了 Linear bottleneck 网络结构。

Linear bottleneck 网络结构中将原来深度可分离卷积中的最后的 ReLU 换成线性激活函数。

现在还有个问题是，深度卷积本身没有改变通道的能力，来的是多少通道输出就是多少通道。如果来的通道很少的话，深度卷积只能在低维度上工作，这样效果并不会很好，所以我们要对通道数进行扩张。既然我们已经知道 PW 逐点卷积也就是 1×1 卷积可以用来升维和降维，那就可以在 DW 深度卷积之前使用 PW 卷积进行升维，再在一个更高维的空间中进行卷积操作来提取特征。

MobileNetV2 网络还有一个改进就是采用像 ResNet 网络那样的残差结构，论文中称为倒残差网络结构。

MobileNet V2 网络主干结构如下图所示：

MobileNet V2 网络组成结构如下图所示。其中 t 表示输入通道的扩张系数； n 表示该模块的重复次数； c 表示输出通道数； s 表示 stride ；注意 stride=2 只在网络结构的第一个第一个卷积层使用，其他的层还是使用 stride=1 。

4.2. 网络实现

定义 ConvBNReLU 网络结构，包括卷积层、标准还、激活函数。注意这里的激活函数为 ReLU6 。当 groups 不是 1 时是深度卷积，为 1 时是普通卷积。

class ConvBNReLU(nn.Sequential):
    def __init__(self, in_channel, out_channel, kernel_size=3, stride=1, groups=1):
        padding = (kernel_size - 1) // 2
        super(ConvBNReLU, self).__init__(
            nn.Conv2d(in_channel, out_channel, kernel_size, stride, padding, groups=groups, bias=False),
            nn.BatchNorm2d(out_channel),
            nn.ReLU6(inplace=True)
        )

定义倒残差网络。

输入参数为：输入通道、输出通道、步长、扩张系数
hidden_channel: 表示增加了扩张系数后的卷积网络的输出
当 stride=1 并且输入通道等于输出通道时使用残差边
如果扩张系数不等于 1 ，先使用 1x1 卷积对输入特征升维
紧接着是深度卷积， groups=hidden_channel
然后在使用 1x1 卷积进行降维
最后使用 BN 层进行归一化
注意：倒残差结构的最后一层没有加激活函数

正向传播时，如果有残差边则将输入与网络的输出相加

class InvertedResidual(nn.Module):
  def __init__(self, in_channel, out_channel, stride, expand_ratio):
      super(InvertedResidual, self).__init__()
      hidden_channel = in_channel * expand_ratio
      self.use_shortcut = stride == 1 and in_channel == out_channel

      layers = []
      if expand_ratio != 1:
          # 1x1 pointwise conv
          layers.append(ConvBNReLU(in_channel, hidden_channel, kernel_size=1))
      layers.extend([
          # 3x3 depthwise conv
          ConvBNReLU(hidden_channel, hidden_channel, stride=stride, groups=hidden_channel),
          # 1x1 pointwise conv(linear)
          nn.Conv2d(hidden_channel, out_channel, kernel_size=1, bias=False),
          nn.BatchNorm2d(out_channel),
      ])

      self.conv = nn.Sequential(*layers)

  def forward(self, x):
      if self.use_shortcut:
          return x + self.conv(x)
      else:
          return self.conv(x)

_make_divisible 函数保证输出的数可以整除 divisor ，原因是在大多数硬件中， size 可以被 d = 8, 16， … 整除的矩阵乘法比较块，因为这些 size 符合处理器单元的对齐位宽。

def _make_divisible(ch, divisor=8, min_ch=None):
    """
    This function is taken from the original tf repo.
    It ensures that all layers have a channel number that is divisible by 8
    It can be seen here:
    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
    """
    if min_ch is None:
        min_ch = divisor
    new_ch = max(min_ch, int(ch + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_ch < 0.9 * ch:
        new_ch += divisor
    return new_ch

MobileNetV2 网络实现代码如下，根据网络结构组成那张图很容易看懂下面的代码。

class MobileNetV2(nn.Module):
    def __init__(self, num_classes=1000, alpha=1.0, round_nearest=8):
        super(MobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel = _make_divisible(32 * alpha, round_nearest)
        last_channel = _make_divisible(1280 * alpha, round_nearest)

        inverted_residual_setting = [
            # t, c, n, s
            [1, 16, 1, 1],
            [6, 24, 2, 2],
            [6, 32, 3, 2],
            [6, 64, 4, 2],
            [6, 96, 3, 1],
            [6, 160, 3, 2],
            [6, 320, 1, 1],
        ]

        features = []
        # conv1 layer
        features.append(ConvBNReLU(3, input_channel, stride=2))
        # building inverted residual residual blockes
        for t, c, n, s in inverted_residual_setting:
            output_channel = _make_divisible(c * alpha, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
                input_channel = output_channel
        # building last several layers
        features.append(ConvBNReLU(input_channel, last_channel, 1))
        # combine feature layers
        self.features = nn.Sequential(*features)

        # building classifier
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(last_channel, num_classes)
        )

        # weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

5. MobileNet V3

5.1. 网络结构

MobileNet V3 相比较 MobileNet V2 引入了轻量级注意力机制以及使用 h-swish 激活函数。

MobileNet V3 包含两个网络： small 和 large ，二者没有明显的区别，只是 bneck 的次数和通道数有一些差异。

small 网络的组成结构如下图所是：

large 网络的组成结构如下图所是：

上述网络组成结构中的定义如下：

第一列 Input 代表 MobileNetV3 每个特征层的输入的 shape ；
第二列 Operator 代表每次特征层即将经历的 block 结构，在 MobileNetV3 中，特征提取经过了许多的 bneck 结构；
第三、四列分别代表了 bneck 内倒残差结构上升后的通道数以及输出通道数。
第五列 SE 代表了是否在这一层引入注意力机制。
第六列 NL 代表了激活函数的种类， HS 代表 h-swish ， RE 代表 RELU 。
第七列 s 代表了每一次 block 结构所用的步长。

h-swish 激活函数公式如下图所是：

MobileNet V3 网络中还引入了注意力机制，网络结构如下所示。

5.2. 网络实现

h-switch 激活函数实现：

class hswish(nn.Module):
    def forward(self, x):
        out = x * F.relu6(x + 3, inplace=True) / 6
        return out

注意力机制实现如下：

首先对输入的特征层进行平均池化，输出大小为 1
然后使用 1x1 卷积，先升维再降维
最后使用 hsigmoid 将输出固定到 [0,1] 之间
再前向传播过程中，将输入乘以注意力机制输出的结果
注意力机制相当与让输入的特征乘以一个权重，让网络更关注权重大的特征 ``` class hsigmoid(nn.Module): def forward(self, x): out = F.relu6(x + 3, inplace=True) / 6 return out

class SeModule(nn.Module): def init(self, in_size, reduction=4): super(SeModule, self).init() self.se = nn.Sequential( nn.AdaptiveAvgPool2d(1), nn.Conv2d(in_size, in_size // reduction, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(in_size // reduction), nn.ReLU(inplace=True), nn.Conv2d(in_size // reduction, in_size, kernel_size=1, stride=1, padding=0, bias=False), nn.BatchNorm2d(in_size), hsigmoid() )

def forward(self, x):
    return x * self.se(x) ```

定义 MobileNet V3 网络结构。

class Block(nn.Module):
    '''expand + depthwise + pointwise'''
    def __init__(self, kernel_size, in_size, expand_size, out_size, nolinear, semodule, stride):
        super(Block, self).__init__()
        self.stride = stride
        self.se = semodule

        self.conv1 = nn.Conv2d(in_size, expand_size, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn1 = nn.BatchNorm2d(expand_size)
        self.nolinear1 = nolinear
        self.conv2 = nn.Conv2d(expand_size, expand_size, kernel_size=kernel_size, stride=stride, padding=kernel_size//2, groups=expand_size, bias=False)
        self.bn2 = nn.BatchNorm2d(expand_size)
        self.nolinear2 = nolinear
        self.conv3 = nn.Conv2d(expand_size, out_size, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn3 = nn.BatchNorm2d(out_size)

        self.shortcut = nn.Sequential()
        if stride == 1 and in_size != out_size:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_size, out_size, kernel_size=1, stride=1, padding=0, bias=False),
                nn.BatchNorm2d(out_size),
            )

    def forward(self, x):
        out = self.nolinear1(self.bn1(self.conv1(x)))
        out = self.nolinear2(self.bn2(self.conv2(out)))
        out = self.bn3(self.conv3(out))
        if self.se != None:
            out = self.se(out)
        out = out + self.shortcut(x) if self.stride==1 else out
        return out

定义 MobileNet V3 small 网络。

class MobileNetV3_Small(nn.Module):
    def __init__(self, num_classes=1000):
        super(MobileNetV3_Small, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.hs1 = hswish()

        self.bneck = nn.Sequential(
            Block(3, 16, 16, 16, nn.ReLU(inplace=True), SeModule(16), 2),
            Block(3, 16, 72, 24, nn.ReLU(inplace=True), None, 2),
            Block(3, 24, 88, 24, nn.ReLU(inplace=True), None, 1),
            Block(5, 24, 96, 40, hswish(), SeModule(40), 2),
            Block(5, 40, 240, 40, hswish(), SeModule(40), 1),
            Block(5, 40, 240, 40, hswish(), SeModule(40), 1),
            Block(5, 40, 120, 48, hswish(), SeModule(48), 1),
            Block(5, 48, 144, 48, hswish(), SeModule(48), 1),
            Block(5, 48, 288, 96, hswish(), SeModule(96), 2),
            Block(5, 96, 576, 96, hswish(), SeModule(96), 1),
            Block(5, 96, 576, 96, hswish(), SeModule(96), 1),
        )

        self.conv2 = nn.Conv2d(96, 576, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(576)
        self.hs2 = hswish()
        self.linear3 = nn.Linear(576, 1280)
        self.bn3 = nn.BatchNorm1d(1280)
        self.hs3 = hswish()
        self.linear4 = nn.Linear(1280, num_classes)
        self.init_params()

    def init_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.hs1(self.bn1(self.conv1(x)))
        out = self.bneck(out)
        out = self.hs2(self.bn2(self.conv2(out)))
        out = F.avg_pool2d(out, 7)
        out = out.view(out.size(0), -1)
        out = self.hs3(self.bn3(self.linear3(out)))
        out = self.linear4(out)
        return out

定义 MobileNet V3 large 网络。

class MobileNetV3_Large(nn.Module):
    def __init__(self, num_classes=1000):
        super(MobileNetV3_Large, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(16)
        self.hs1 = hswish()

        self.bneck = nn.Sequential(
            Block(3, 16, 16, 16, nn.ReLU(inplace=True), None, 1),
            Block(3, 16, 64, 24, nn.ReLU(inplace=True), None, 2),
            Block(3, 24, 72, 24, nn.ReLU(inplace=True), None, 1),
            Block(5, 24, 72, 40, nn.ReLU(inplace=True), SeModule(40), 2),
            Block(5, 40, 120, 40, nn.ReLU(inplace=True), SeModule(40), 1),
            Block(5, 40, 120, 40, nn.ReLU(inplace=True), SeModule(40), 1),
            Block(3, 40, 240, 80, hswish(), None, 2),
            Block(3, 80, 200, 80, hswish(), None, 1),
            Block(3, 80, 184, 80, hswish(), None, 1),
            Block(3, 80, 184, 80, hswish(), None, 1),
            Block(3, 80, 480, 112, hswish(), SeModule(112), 1),
            Block(3, 112, 672, 112, hswish(), SeModule(112), 1),
            Block(5, 112, 672, 160, hswish(), SeModule(160), 1),
            Block(5, 160, 672, 160, hswish(), SeModule(160), 2),
            Block(5, 160, 960, 160, hswish(), SeModule(160), 1),
        )

        self.conv2 = nn.Conv2d(160, 960, kernel_size=1, stride=1, padding=0, bias=False)
        self.bn2 = nn.BatchNorm2d(960)
        self.hs2 = hswish()
        self.linear3 = nn.Linear(960, 1280)
        self.bn3 = nn.BatchNorm1d(1280)
        self.hs3 = hswish()
        self.linear4 = nn.Linear(1280, num_classes)
        self.init_params()

    def init_params(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                init.constant_(m.weight, 1)
                init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                init.normal_(m.weight, std=0.001)
                if m.bias is not None:
                    init.constant_(m.bias, 0)

    def forward(self, x):
        out = self.hs1(self.bn1(self.conv1(x)))
        out = self.bneck(out)
        out = self.hs2(self.bn2(self.conv2(out)))
        out = F.avg_pool2d(out, 7)
        out = out.view(out.size(0), -1)
        out = self.hs3(self.bn3(self.linear3(out)))
        out = self.linear4(out)
        return out

YOLO 源代码分析(四) YOLOv3 网络训练

Thu, 02 Jul 2020 00:00:00 +0000

这篇文章介绍 YOLOv3 的网络训练过程。主要介绍网络的正负样本的选取和损失函数计算。

训练模型的代码在 train.py 文件中。这个文件大部分内容前面三篇文章都介绍过了，这里只讲解之前每提到的内容。

下面的 accumulate 是模拟一个更大的 batchsize 来进行梯度下降，一定条件下， batchsize 越大训练效果越好，梯度累加则实现了 batchsize 的变相扩大，如果 accumulate 为 8 ，则 batchsize ‘变相’ 扩大了 8 倍。

accumulate = max(round(64 / batch_size), 1)

下面代码是使用多尺度训练。

gs: 首先确保网络输入的图像是 32 的倍数

multi_scale: 多尺度训练，就是在训练过程中不断的改变训练数据集的大小，增加模型的鲁棒性，每训练 accumulate 个 batch 后，修改图像训练图像的大小，图像大小在 (grid_min, grid_max) 之间，并且都是 32 的倍数

gs = 32  # (pixels) grid size
assert math.fmod(imgsz_test, gs) == 0, "--img-size %g must be a %g-multiple" % (imgsz_test, gs)
grid_min, grid_max = imgsz_test // gs, imgsz_test // gs
if multi_scale:
  imgsz_min = opt.img_size // 1.5
  imgsz_max = opt.img_size // 0.667

  grid_min, grid_max = imgsz_min // gs, imgsz_max // gs
  imgsz_min, imgsz_max = int(grid_min * gs), int(grid_max * gs)
  imgsz_train = imgsz_max  # initialize with max size
  print("Using multi_scale training, image range[{}, {}]".format(imgsz_min, imgsz_max))

然后我们直接进入 train_one_epoch 函数，这个函数用来训练一个 epoch 。

第一个 epoch 使用 Warmup 训练模式。为什么要使用 Warmup 训练模式呢？由于刚开始训练时，模型的权重 (weights) 是随机初始化的，此时若选择一个较大的学习率，可能带来模型的不稳定，选择 Warmup 预热学习率的方式，可以使得开始训练的几个 epoch 或者一些 step 内学习率较小，在预热的小学习率下，模型可以慢慢趋于稳定，等模型相对稳定后在选择预先设置的学习率进行训练，使得模型收敛速度变得更快，模型效果更佳。

if epoch == 0 and warmup is True:  # 当训练第一轮（epoch=0）时，启用warmup训练方式，可理解为热身训练
    warmup_factor = 1.0 / 1000
    warmup_iters = min(1000, len(data_loader) - 1)

    lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)
    accumulate = 1

每训练 accumulate 个 batch 就随机修改一次输入图片大小，由于 label 已转为相对坐标，故缩放图片不影响 label 的值。

img_size: 在给定最大最小输入尺寸范围内随机选取一个 size (size 为 32 的整数倍)

如果图片最大边长不等于 img_size , 则缩放图片，并将长和宽调整到 32 的整数倍

if multi_scale:
  if ni % accumulate == 0:  #  adjust img_size (67% - 150%) every 1 batch
      img_size = random.randrange(grid_min, grid_max + 1) * gs
  sf = img_size / max(imgs.shape[2:])  # scale factor

  if sf != 1:
      # gs: (pixels) grid size
      ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to 32-multiple)
      imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)

每训练 accumulate 个 batch 更新一次权重。

if ni % accumulate == 0:
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

下面看下损失函数的计算。先看 build_targets 函数，该函数用来划分正负样本。

nt: 记录当前 batch 中有多少个目标，目标的 shape 是： (image,class,x,y,w,h)
tcls, tbox, indices, anch: 最后解释
gain: 特征图的大小
multi_gpu: 是否使用多 GPU
遍历每一个 YOLO 层
anchors: 获取该 yolo predictor 对应的 anchors
na: anchor 的个数，每个 yolo 层有 3 个大小的 anchor
at: 对 anchor 进行维度上的变换，变换过程如下： [3] -> [3, 1] -> [3, nt] ，比如 nt=14 表示当前 batch 有 14 个目标， at 转化成 (3,14) 的 tensor ，表示 3 个 anchor 和 14 个目标
接下来是为每一个目标匹配 anchor
t=targets * gain ， targets 为目标边界框，被缩放到图片的相对尺寸下， gain 为 feature map 的大小， targets * gain 就是把目标边界框映射到 featrue map 上
wh_iou: 计算 anchor 与目标的 iou ，这里计算 IOU 的方法不同于以往我见到过的，以往都是传入左上角和右下角两个坐标，而 wh_iou 只传入了宽和高。这里我查了很多资料，也思考了很久，说一下我的理解：训练时，目标的中心点坐标落在哪个 cell 里，哪个 cell 就负责预测这个目标， yolov3 有三个预测特征层，分别预测不同大小的目标，每个预测特征层的每个 cell 有 3 个 anchor ，那到底用哪个 anchor 进行预测呢？这句代码就是为一个 batch 中的每个目标选取合适的 anchor ，这里的合适就是将 anchor 和目标的左上角对齐计算 iou ，大于阈值(0.2) 则让这个 anchor 负责预测这个目标。为什么这样选？我的理解是， anchor 是在固定位置设置的候选框，需要通过网络预测的边界框参数进行调整才能得到最终的预测框。只有 anchor 和目标的 iou 大于设定的这个阈值时，才有可能通过预测的边界框回归参数将其调整到目标的大小。这里在进行筛选的时候，会出现一个目标由多个 anchor 进行预测的情况，这个问题在测试阶段非极大值抑制会过滤掉重复的目标，在训练阶段，网络只管让负责预测这个目标的 anchor 尽可能的接近目标。
j: 保存 anchor 与目标的 iou 大于 model.hyp[‘iou_t’] 的 bool 索引，shape 为 [3, nt] ，其中为 True 的是符合的 anchor 。
a, t = at[j], t.repeat(na, 1, 1)[j]: 获取 iou 大于阈值的 anchor 与 target 对应信息。t.repeat(na, 1, 1): [nt, 6] -> [3, nt, 6]
b, c: b 表示图片在 batch中的索引， c 表示类别
gxy: 获取目标的中心点坐标
gij: 获取目标中心点坐标落在哪个 cell 上， gjj 表示 cell 的左上角坐标
gi, gj: 获取 cell 的左上角坐标
indices: 保存 image 、 anchor 、 grid 的索引
tbox: 相对于 cell 左上角的偏移和 gwh
anch: 保存 anchor

tcls: 保存 class

def build_targets(p, targets, model):
  nt = targets.shape[0]
  tcls, tbox, indices, anch = [], [], [], []
  gain = torch.ones(6, device=targets.device)  # normalized to gridspace gain

  multi_gpu = type(model) in (nn.parallel.DataParallel, nn.parallel.DistributedDataParallel)
  for i, j in enumerate(model.yolo_layers):  # [89, 101, 113]
      anchors = model.module.module_list[j].anchor_vec if multi_gpu else model.module_list[j].anchor_vec
      gain[2:] = torch.tensor(p[i].shape)[[3, 2, 3, 2]]  # xyxy gain
      na = anchors.shape[0]  # number of anchors
      at = torch.arange(na).view(na, 1).repeat(1, nt)  # anchor tensor, same as .repeat_interleave(nt)

      # Match targets to anchors
      a, t, offsets = [], targets * gain, 0
      if nt:  # 如果存在target的话
          j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t']
          a, t = at[j], t.repeat(na, 1, 1)[j]  # filter

      b, c = t[:, :2].long().T  # image, class
      gxy = t[:, 2:4]  # grid xy
      gwh = t[:, 4:6]  # grid wh
      gij = (gxy - offsets).long()  # 匹配targets所在的grid cell左上角坐标
      gi, gj = gij.T  # grid xy indices

      # Append
      indices.append((b, a, gj, gi))  # image, anchor, grid indices(x, y)
      tbox.append(torch.cat((gxy - gij, gwh), 1))  # gt box相对anchor的x,y偏移量以及w,h
      anch.append(anchors[a])  # anchors
      tcls.append(c)  # class
      if c.shape[0]:  # if any targets
          # 目标的标签数值不能大于给定的目标类别数
          assert c.max() < model.nc, 'Model accepts %g classes labeled from 0-%g, however you labelled a class %g. ' \
                                     'See https://github.com/ultralytics/yolov3/wiki/Train-Custom-Data' % (
                                         model.nc, model.nc - 1, c.max())

  return tcls, tbox, indices, anch

损失函数。

类别和置信度损失使用二值交叉熵损失函数
smooth_BCE: 下面介绍
FocalLoss: 下面介绍
定位损失使用 GIou 损失函数
p 存储 yolo 的输出，遍历每一个 yolo 的输出层， shape 是： [bs, anchor, grid, grid, xywh + obj + classes]
ps = pi[b, a, gj, gi]: 对应匹配到正样本的预测信息
用匹配到的预测信息与真实标签计算损失，先计算 GIou 损失
tobj: 置信度损失，当前框有目标的概率乘以 bounding box 和 ground truth 的 IoU 的结果
lcls: 类别损失

最后将各类损失乘以各自的权重

def compute_loss(p, targets, model):  # predictions, targets, model
  device = p[0].device
  lcls = torch.zeros(1, device=device)  # Tensor(0)
  lbox = torch.zeros(1, device=device)  # Tensor(0)
  lobj = torch.zeros(1, device=device)  # Tensor(0)
  tcls, tbox, indices, anchors = build_targets(p, targets, model)  # targets
  h = model.hyp  # hyperparameters
  red = 'mean'  # Loss reduction (sum or mean)

  # Define criteria
  BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device), reduction=red)
  BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device), reduction=red)

  # class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3
  cp, cn = smooth_BCE(eps=0.0)

  # focal loss
  g = h['fl_gamma']  # focal loss gamma
  if g > 0:
      BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g)

  # per output
  nt = 0  # targets
  for i, pi in enumerate(p):  # layer index, layer predictions
      b, a, gj, gi = indices[i]  # image, anchor, gridy, gridx
      tobj = torch.zeros_like(pi[..., 0], device=device)  # target obj

      nb = b.shape[0]  # number of targets
      if nb:
          nt += nb  # cumulative targets
          ps = pi[b, a, gj, gi]  # prediction subset corresponding to targets

          # GIoU
          pxy = ps[:, :2].sigmoid()
          pwh = ps[:, 2:4].exp().clamp(max=1E3) * anchors[i]
          pbox = torch.cat((pxy, pwh), 1)  # predicted box
          giou = bbox_iou(pbox.t(), tbox[i], x1y1x2y2=False, GIoU=True)  # giou(prediction, target)
          lbox += (1.0 - giou).mean()  # giou loss

          # Obj
          tobj[b, a, gj, gi] = (1.0 - model.gr) + model.gr * giou.detach().clamp(0).type(tobj.dtype)  # giou ratio

          # Class
          if model.nc > 1:  # cls loss (only if multiple classes)
              t = torch.full_like(ps[:, 5:], cn, device=device)  # targets
              t[range(nb), tcls[i]] = cp
              lcls += BCEcls(ps[:, 5:], t)  # BCE

      lobj += BCEobj(pi[..., 4], tobj)  # obj loss

  lbox *= h['giou']
  lobj *= h['obj']
  lcls *= h['cls']

  # loss = lbox + lobj + lcls
  return {"box_loss": lbox, "obj_loss": lobj, "class_loss": lcls}

下面这段代码提供了 label smoothing 的功能，只能用在多类问题， ultralytics 版 YOLOv3 使用的二分类进行损失计算，所以实际使用中代码中传入的参数为 0 。

def smooth_BCE(eps=0.1):
    # return positive, negative label smoothing BCE targets
    return 1.0 - 0.5 * eps, 0.5 * eps

这里介绍下为什么要引入 label smoothing 。交叉熵损失函数在多分类任务中存在如下问题，训练神经网络时，最小化预测概率和标签真实概率之间的交叉熵，从而得到最优的预测概率分布。神经网络会促使自身往正确标签和错误标签差值最大的方向学习，在训练数据较少，不足以表征所有的样本特征的情况下，会导致网络过拟合。 label smoothing 可以解决上述问题，这是一种正则化策略，主要是通过 soft one-hot 来加入噪声，减少了真实样本标签的类别在计算损失函数时的权重，最终起到抑制过拟合的效果。

下面看下 FocalLoss 的实现。通过增加 gamma 和 alpha 两个参数提高模型对难分和易分样本极度不平衡问题。

reduction: 该参数共有三种选项 mean，sum 和 none 。 mean 为默认情况，表明对 N 个样本的 loss 进行求平均之后返回； sum 指对 N 个样本的 loss 求和； none 表示直接返回 N 分样本的 loss 。

class FocalLoss(nn.Module):
  # Wraps focal loss around existing loss_fcn(), i.e. criteria = FocalLoss(nn.BCEWithLogitsLoss(), gamma=1.5)
  def __init__(self, loss_fcn, gamma=1.5, alpha=0.25):
      super(FocalLoss, self).__init__()
      self.loss_fcn = loss_fcn  # must be nn.BCEWithLogitsLoss()
      self.gamma = gamma
      self.alpha = alpha
      self.reduction = loss_fcn.reduction
      self.loss_fcn.reduction = 'none'  # required to apply FL to each element

  def forward(self, pred, true):
      loss = self.loss_fcn(pred, true)

      pred_prob = torch.sigmoid(pred)  # prob from logits
      p_t = true * pred_prob + (1 - true) * (1 - pred_prob)
      alpha_factor = true * self.alpha + (1 - true) * (1 - self.alpha)
      modulating_factor = (1.0 - p_t) ** self.gamma
      loss *= alpha_factor * modulating_factor

      if self.reduction == 'mean':
          return loss.mean()
      elif self.reduction == 'sum':
          return loss.sum()
      else:  # 'none'
          return loss

YOLO 源代码分析(三) YOLOv3 主干网络

Thu, 25 Jun 2020 00:00:00 +0000

1. 概述

ultralytics 版 YOLOv3 的主干网络是由 Darknet53 加上 SPP 结构再加上残差网络构成的。这个版本的主干网络是通过解析配置文件来加载的，由于网络整体结构比较简单，不过多介绍，本文主要介绍如果解析配置文件。

2. 配置文件格式

主干网络的配置文件再 my_yolov3.cfg 配置文件中，在搭建自己的网络时只需要根据网络的预测类别进行相关参数的修改即可。如果使用 VOC 数据集，数据集的类别是 20 ，每个预测器预测 (20 + 1 + 4)=25 个值，分别为 20 个类别和背景，再加上边界框回归参数， feature map 上的每一个点预测 3 个 anchor 也就是每个点预测 75 个值。

2.1. 类型

net 该类型定义网络的超参数。

convolutional 该类型定义网络卷积神经网络。

shortcut 残差网络，各个元素对应相加。

maxpool 最大池化层。

route 残差网络，在 channel 维度上进行拼接。

yolo 网络预测层。

upsample 上采样层。

2.2. 实现

首先看 Darknet 网络的实现，构造函数实现如下：

parse_model_cfg: 解析配置文件，以字典的形式缓存配置
create_modules: 根据配置文件构建网络，下面重点介绍

self.yolo_layers: 获取网络中的 yolo 层，一共有 3 个 yolo 层

class Darknet(nn.Module):
  def __init__(self, cfg, img_size=(416, 416), verbose=False):
      super(Darknet, self).__init__()
      self.input_size = [img_size] * 2 if isinstance(img_size, int) else img_size
      self.module_defs = parse_model_cfg(cfg)
      self.module_list, self.routs = create_modules(self.module_defs, img_size)
      self.yolo_layers = get_yolo_layers(self)
      self.info(verbose) if not ONNX_EXPORT else None  # print model description

让我们看一下网络中的各个模块是如何搭建起来的，在 create_modules 函数中，根据配置文的缓存数据进行搭建网络：

从缓存中删除列表的第一项，该项为网络超参数，未使用，超参数在 hyp.yaml 中
output_filters: 该变量记录每一层的输入通道数，初始值为 3 ，输入的是 RGB 彩色图像
module_list: 该变量存储网络的每一个模块
routs: 改变量统计哪些特征层的输出会被后续的层使用到(可能是特征融合，也可能是拼接)
yolo_index: 存储 yolo 层的索引，一共有 3 个 yolo 层
接下来是遍历 modules_defs 中的每一项，存储的是一个字典
convolutional: 卷积层，添加 Conv2d 模块，添加 activation 激活函数，添加 BatchNorm2d ，如果不存在 BN 层，意味着该层为 yolo 层，将层号添加到 routs 变量中
maxpool: 最大池化层
upsample: 上采样层
route: 这一层代表残差网络，在 channel 维度上进行拼接。 layers 表示与哪些层进行拼接， layers 是负值，表示与 output_filters 倒数第 layers 层进行拼接， layers 为正值，表示与 output_filters 正数第 layers+1 层为与之拼接的层进行拼接； filters 求和就是经过拼接后网络输出的维度，添加到 output_filter 中； routs 记录层号，这些层的输出以后会用到，当 l < 0 时，向前跳过 l 层，为拼接的层，当 l > 0 时，则第 l 层为拼接的层
shortcut: 这一层表示特征融合，对应元素相加。融合后网络的维度， shortcut 是对应元素相加，维度不变。
yolo: 网络有 3 个，每个 yolo 层的缩放比例分别为 [32, 16, 8] ，下面对于 bias_ 的设置我目前还每看懂是什么意思。

最后返回 module_list 和 routs_binary ， routs_binary 记录哪些层需要融合或者拼接

def create_modules(modules_defs: list, img_size):
  img_size = [img_size] * 2 if isinstance(img_size, int) else img_size
  modules_defs.pop(0)  # cfg training hyperparams (unused)
  output_filters = [3]  # input channels
  module_list = nn.ModuleList()
  routs = []  # list of layers which rout to deeper layers
  yolo_index = -1

  for i, mdef in enumerate(modules_defs):
      modules = nn.Sequential()

      if mdef["type"] == "convolutional":
          bn = mdef["batch_normalize"]  # 1 or 0 / use or not
          filters = mdef["filters"]
          k = mdef["size"]  # kernel size
          stride = mdef["stride"] if "stride" in mdef else (mdef['stride_y'], mdef["stride_x"])
          if isinstance(k, int):
              modules.add_module("Conv2d", nn.Conv2d(in_channels=output_filters[-1],
                                                     out_channels=filters,
                                                     kernel_size=k,
                                                     stride=stride,
                                                     padding=k // 2 if mdef["pad"] else 0,
                                                     bias=not bn))
          else:
              raise TypeError("conv2d filter size must be int type.")

          if bn:
              modules.add_module("BatchNorm2d", nn.BatchNorm2d(filters))
          else:
              routs.append(i)  # detection output (goes into yolo layer)

          if mdef["activation"] == "leaky":
              modules.add_module("activation", nn.LeakyReLU(0.1, inplace=True))
          else:
              pass

      elif mdef["type"] == "BatchNorm2d":
          pass

      elif mdef["type"] == "maxpool":
          k = mdef["size"]  # kernel size
          stride = mdef["stride"]
          modules = nn.MaxPool2d(kernel_size=k, stride=stride, padding=(k - 1) // 2)

      elif mdef["type"] == "upsample":
          if ONNX_EXPORT:  # explicitly state size, avoid scale_factor
              g = (yolo_index + 1) * 2 / 32  # gain
              modules = nn.Upsample(size=tuple(int(x * g) for x in img_size))
          else:
              modules = nn.Upsample(scale_factor=mdef["stride"])

      elif mdef["type"] == "route":  # [-2],  [-1,-3,-5,-6], [-1, 61]
          layers = mdef["layers"]
          filters = sum([output_filters[l + 1 if l > 0 else l] for l in layers])
          routs.extend([i + l if l < 0 else l for l in layers])
          modules = FeatureConcat(layers=layers)

      elif mdef["type"] == "shortcut":
          layers = mdef["from"]
          filters = output_filters[-1]
          routs.append(i + layers[0])
          modules = WeightedFeatureFusion(layers=layers, weight="weights_type" in mdef)

      elif mdef["type"] == "yolo":
          yolo_index += 1  # 记录是第几个yolo_layer [0, 1, 2]
          stride = [32, 16, 8]  # 预测特征层对应原图的缩放比例

          modules = YOLOLayer(anchors=mdef["anchors"][mdef["mask"]],  # anchor list
                              nc=mdef["classes"],  # number of classes
                              img_size=img_size,
                              stride=stride[yolo_index])

          # Initialize preceding Conv2d() bias (https://arxiv.org/pdf/1708.02002.pdf section 3.3)
          try:
              j = -1
              bias_ = module_list[j][0].bias  # shape(255,) 索引0对应Sequential中的Conv2d
              bias = bias_.view(modules.na, -1)  # shape(3, 85)
              bias[:, 4] += -4.5  # obj
              bias[:, 5:] += math.log(0.6 / (modules.nc - 0.99))  # cls (sigmoid(p) = 1/nc)
              module_list[j][0].bias = torch.nn.Parameter(bias_, requires_grad=bias_.requires_grad)
          except Exception as e:
              print('WARNING: smart bias initialization failure.', e)
      else:
          print("Warning: Unrecognized Layer Type: " + mdef["type"])

      # Register module list and number of output filters
      module_list.append(modules)
      output_filters.append(filters)

  routs_binary = [False] * len(modules_defs)
  for i in routs:
      routs_binary[i] = True
  return module_list, routs_binary

接下来介绍 route 层的实现，该层实现在 FeatureConcat 类中：

将多个特征层在 channel 维度上拼接在一起

class FeatureConcat(nn.Module):
  def __init__(self, layers):
      super(FeatureConcat, self).__init__()
      self.layers = layers  # layer indices
      self.multiple = len(layers) > 1  # multiple layers flag

  def forward(self, x, outputs):
      return torch.cat([outputs[i] for i in self.layers], 1) if self.multiple else outputs[self.layers[0]]

特征融合的实现在 WeightedFeatureFusion 类中。

默认 weight=False ，所以前向传播就是直接各个层的输出特征相加

class WeightedFeatureFusion(nn.Module):  # weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070
  def __init__(self, layers, weight=False):
      super(WeightedFeatureFusion, self).__init__()
      self.layers = layers  # layer indices
      self.weight = weight  # apply weights boolean
      self.n = len(layers) + 1  # number of layers 融合的特征矩阵个数
      if weight:
          self.w = nn.Parameter(torch.zeros(self.n), requires_grad=True)  # layer weights

  def forward(self, x, outputs):
      # Weights
      if self.weight:
          w = torch.sigmoid(self.w) * (2 / self.n)  # sigmoid weights (0-1)
          x = x * w[0]

      # Fusion
      nx = x.shape[1]  # input channels
      for i in range(self.n - 1):
          a = outputs[self.layers[i]] * w[i + 1] if self.weight else outputs[self.layers[i]]  # feature to add
          na = a.shape[1]  # feature channels

          if nx == na:  # same shape 如果channel相同，直接相加
              x = x + a
          elif nx > na:  # slice input 如果channel不同，将channel多的特征矩阵砍掉部分channel保证相加的channel一致
              x[:, :na] = x[:, :na] + a  # or a = nn.ZeroPad2d((0, 0, 0, 0, 0, dc))(a); x = x + a
          else:  # slice feature
              x = x + a[:, :nx]

      return x

预测器实现在 YOLOLayer 类中，构造函数实现如下：

传入当前预测特征层的 anchor 的大小，类别个数， img_size ，下采样倍数
self.anchor_vec: 将 anchors 大小缩放到 grid 尺度

self.anchor_wh: 各个元素的定义为 [batch_size, na, grid_h, grid_w, wh] ，值为 1 的维度对应的值不是固定值，后续操作可根据 broadcast 广播机制自动扩充

class YOLOLayer(nn.Module):
  def __init__(self, anchors, nc, img_size, stride):
      super(YOLOLayer, self).__init__()
      self.anchors = torch.Tensor(anchors)
      self.stride = stride
      self.na = len(anchors)  # number of anchors (3)
      self.nc = nc  # number of classes (80)
      self.no = nc + 5  # number of outputs (85: x, y, w, h, obj, cls1, ...)
      self.nx, self.ny, self.ng = 0, 0, (0, 0)  # initialize number of x, y gridpoints
      self.anchor_vec = self.anchors / self.stride
      self.anchor_wh = self.anchor_vec.view(1, self.na, 1, 1, 2)
      self.grid = None

      if ONNX_EXPORT:
          self.training = False
          self.create_grids((img_size[1] // stride, img_size[0] // stride))  # number x, y grid points

接下来看 forward 函数，首先是 create_grids 函数，该函数主要用于生成网格，且在推理时才使用。

self.ng: 生成网格大小
使用 torch.meshgrid 函数在特征图上获得网格点的坐标

self.grid: 将网格点坐标堆叠在一起，将坐标 reshape 到 (batch_size, na, grid_h, grid_w, wh) 维度，值为 1 的可根据广播机制进行扩充

def create_grids(self, ng=(13, 13), device="cpu"):
  self.nx, self.ny = ng
  self.ng = torch.tensor(ng, dtype=torch.float)

  # build xy offsets 构建每个cell处的anchor的xy偏移量(在feature map上的)
  if not self.training:  # 训练模式不需要回归到最终预测boxes
      yv, xv = torch.meshgrid([torch.arange(self.ny, device=device),
                                  torch.arange(self.nx, device=device)])
      # batch_size, na, grid_h, grid_w, wh
      self.grid = torch.stack((xv, yv), 2).view((1, 1, self.ny, self.nx, 2)).float()

  if self.anchor_vec.device != device:
      self.anchor_vec = self.anchor_vec.to(device)
      self.anchor_wh = self.anchor_wh.to(device)

下面看 forward 函数的实现：

对输入进行维度的调整 [bs, anchor, grid, grid, xywh + obj + classes]
如果时训练模式，直接返回调整维度后的值，进行损失计算
如果时验证模式，将网络输出值处理，首先将 xy 转换成在 feature map 上的 xy 坐标，将 wh 转换成在 feature map 上的 wh 值，将 xywh 四个值映射到原图上，计算网络输出的置信度。最后将预测值 reshape 到 [bs, -1, self.no] 格式。 ``` def forward(self, p): if ONNX_EXPORT: bs = 1 # batch size else: bs, _, ny, nx = p.shape # batch_size, predict_param(255), grid(13), grid(13) if (self.nx, self.ny) != (nx, ny) or self.grid is None: # fix no grid bug self.create_grids((nx, ny), p.device)

p = p.view(bs, self.na, self.no, self.ny, self.nx).permute(0, 1, 3, 4, 2).contiguous() # prediction

if self.training: return p elif ONNX_EXPORT: # Avoid broadcasting for ANE operations m = self.na * self.nx * self.ny # 3* ng = 1. / self.ng.repeat(m, 1) grid = self.grid.repeat(1, self.na, 1, 1, 1).view(m, 2) anchor_wh = self.anchor_wh.repeat(1, 1, self.nx, self.ny, 1).view(m, 2) * ng
```
  p = p.view(m, self.no)
  # xy = torch.sigmoid(p[:, 0:2]) + grid  # x, y
  # wh = torch.exp(p[:, 2:4]) * anchor_wh  # width, height
  # p_cls = torch.sigmoid(p[:, 4:5]) if self.nc == 1 else \
  #     torch.sigmoid(p[:, 5:self.no]) * torch.sigmoid(p[:, 4:5])  # conf
  p[:, :2] = (torch.sigmoid(p[:, 0:2]) + grid) * ng  # x, y
  p[:, 2:4] = torch.exp(p[:, 2:4]) * anchor_wh  # width, height
  p[:, 4:] = torch.sigmoid(p[:, 4:])
  p[:, 5:] = p[:, 5:self.no] * p[:, 4:5]
  return p   else:  # inference
  # [bs, anchor, grid, grid, xywh + obj + classes]
  io = p.clone()  # inference output
  io[..., :2] = torch.sigmoid(io[..., :2]) + self.grid
  io[..., 2:4] = torch.exp(io[..., 2:4]) * self.anchor_wh
  io[..., :4] *= self.stride
  torch.sigmoid_(io[..., 4:])
  return io.view(bs, -1, self.no), p  # view [1, 3, 13, 13, 85] as [1, 507, 85]
```

&#160; &#160; &#160; &#160;Darknet 网络中的子模块就介绍完了，现在我们回到 Darknet 实现上。在 Darknet 构造函数中调用了获取 yolo 层的接口，实现如下：
* yolo 层共有 3 个，分别是 [89, 101, 113]

def get_yolo_layers(self): return [i for i, m in enumerate(self.module_list) if m.class.name == ‘YOLOLayer’]

&#160; &#160; &#160; &#160;下面分析一下 Darknet 的前向传播函数， forward 函数直接调用了 forward_once 函数。
* yolo_out, out: yolo_out 收集每个 yolo_layer 层的输出； out 收集每个模块的输出。
* 前向传播过程中要记录需要网络的输出，以供 WeightedFeatureFusion 和 FeatureConcat 两个层使用。
* 如果是训练模式，返回 yolo 层的输出， shape 为 [bs, anchor, grid, grid, xywh + obj + classes]
* 如果是验证模式，返回 yolo 层的输出， shape 为 [bs, -1, self.no] 和 [bs, anchor, grid, grid, xywh + obj + classes] ，最后要将 [bs, -1, self.no] 在维度 1 上进行堆叠。

def forward(self, x, verbose=False): return self.forward_once(x, verbose=verbose)

def forward_once(self, x, verbose=False): yolo_out, out = [], []

for i, module in enumerate(self.module_list):
    name = module.__class__.__name__
    if name in ["WeightedFeatureFusion", "FeatureConcat"]:  # sum, concat
        x = module(x, out)  # WeightedFeatureFusion(), FeatureConcat()
    elif name == "YOLOLayer":
        yolo_out.append(module(x))
    else:  # run module directly, i.e. mtype = 'convolutional', 'upsample', 'maxpool', 'batchnorm2d' etc.
        x = module(x)

    out.append(x if self.routs[i] else [])

if self.training:  # train
    return yolo_out
elif ONNX_EXPORT:  # export
    p = torch.cat(yolo_out, dim=0)

    return p
else:  # inference or test
    x, p = zip(*yolo_out)  # inference output, training output
    x = torch.cat(x, 1)  # cat yolo outputs

    return x, p ``` &#160; &#160; &#160; &#160;至此， YOLOv3 主干网络的源码分析就结束了。