Gaze Estimation using Transformer代码模块解读

代码链接：GitHub - yihuacheng/GazeTR: The codes and models in 'Gaze Estimation using Transformer, ICPR2022'.

 项目主目录为GazeTR，在主目录下有许多子模块。以下为分模块解读。

config

config模块下有test和train俩个子模块，其中均为.yaml文件，分别是在四个gaze数据集下，对一些超参数的基本设置，以及数据集的存储路径。以训练gaze360数据集为例，相关设置如下：

 reader

reader模块定义了一个用于读取数据的PyTorch数据加载器，用于加载和预处理训练用的图像和标签数据。

Decode_MPII、Decode_Diap、Decode_Gaze360、Decode_ETH和Decode_RTGene函数用于解析不同数据集的注释。它们将注释数据的一行作为输入，并返回一个包含相关字段（例如，面部图像文件名、2D凝视坐标等）作为键的字典对象。

 Get_Decode函数将数据集名称作为输入，并返回相应的注释解码函数。

def Get_Decode(name):
    mapping = Decode_Dict()
    keys = list(mapping.keys())
    name = name.lower()
    score = [long_substr(name, i) for i in keys]
    key  = keys[score.index(max(score))]
    return mapping[key]

trainloader类读取源数据（图像文件名和注释文件）并将其存储在一个edict对象中。它还定义了一个transforms对象，用于将预处理操作应用于输入图像。

trainloader类的__len__方法返回数据集中样本的数量。

trainloader类的__getitem__方法以索引值作为输入，并返回相应的预处理图像和标签数据。

loader函数使用指定的源数据和批量大小初始化trainloader对象，并返回一个DataLoader对象，该对象可用于在训练期间迭代数据批次。

def loader(source, batch_size, shuffle=True,  num_workers=0):
    dataset = trainloader(source)
    print(f"-- [Read Data]: Source: {source.label}")
    print(f"-- [Read Data]: Total num: {len(dataset)}")
    load = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)
    return load

tester&trainer

这俩大模块是项目的测试模块和训练模块。

model

除上面这些模块以外，作者定义了自己设计的模型函数

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        maps = 32
        nhead = 8
        dim_feature = 7*7
        dim_feedforward=512
        dropout = 0.1
        num_layers=6

        self.base_model = resnet18(pretrained=False, maps=maps)

        # d_model: dim of Q, K, V
        # nhead: seq num
        # dim_feedforward: dim of hidden linear layers
        # dropout: prob

        encoder_layer = TransformerEncoderLayer(
                  maps,
                  nhead,
                  dim_feedforward,
                  dropout)

        encoder_norm = nn.LayerNorm(maps)
        # num_encoder_layer: deeps of layers

        self.encoder = TransformerEncoder(encoder_layer, num_layers, encoder_norm)

        self.cls_token = nn.Parameter(torch.randn(1, 1, maps))

        self.pos_embedding = nn.Embedding(dim_feature+1, maps)

        self.feed = nn.Linear(maps, 2)

        self.loss_op = nn.L1Loss()


    def forward(self, x_in):
        feature = self.base_model(x_in["face"])
        batch_size = feature.size(0)
        feature = feature.flatten(2)
        feature = feature.permute(2, 0, 1)
        
        cls = self.cls_token.repeat( (1, batch_size, 1))
        feature = torch.cat([cls, feature], 0)
        
        position = torch.from_numpy(np.arange(0, 50)).cuda()

        pos_feature = self.pos_embedding(position)

        # feature is [HW, batch, channel]
        feature = self.encoder(feature, pos_feature)
  
        feature = feature.permute(1, 2, 0)

        feature = feature[:,:,0]

        gaze = self.feed(feature)
        
        return gaze

    def loss(self, x_in, label):
        gaze = self.forward(x_in)
        loss = self.loss_op(gaze, label) 
        return loss

使用ResNet-18的卷积层进行特征提取。卷积层从人脸图像生成7 × 7 × 512个特征图。然后使用额外的1 × 1卷积层来缩放通道，并得到7 × 7 × 32个特征图。将特征映射输入到一个6层transformer中。对于transformer，将两层MLP的隐藏大小设置为512，并执行8头自注意机制。dropout设置为0.1。