Model Architecture and Loss Functions

def Conv_Block(in_channel, out_channel, kernel_size, stride, padding, group=1, has_bn=True, is_linear=False):
    '''
    Conv_Block 的 Docstring

    :param in_channel: 输入通道数
    :param out_channel: 输出通道数
    :param kernel_size: 卷积核尺寸，类型是 int 或则和 tuple
    :param stride: 卷积的步长，类型是 int
    :param padding: 边缘填充的大小，用于确保 feature map 尺寸不会缩小，与 stride 配合使用，类型是 int
    :param group: groups = group，默认为1，表示标准卷积；groups = in_channel，则表示深度可分离卷积
    :param has_bn: 是否使用批归一化，类型是 bool
    :param is_linear: 是否使用线性激活函数，类型是 bool
    '''
    return Sequential(
        Conv2d(in_channel, out_channel, kernel_size, stride, padding=padding, groups=group, bias=False),
        BatchNorm2d(out_channel) if has_bn else Sequential(),
        ReLU(inplace=True) if not is_linear else Sequential()
    )

class InvertedResidual(Module):
    def __init__(self, in_channel, out_channel, stride, use_res_connect, expand_ratio):
        super(InvertedResidual, self).__init__()
        # 调用父类构造函数，这样才能使用.cuda(), .parameters(), .state_dict()等方法以及forward魔术方法
        self.stride = stride
        assert stride in [1, 2]
        # 含义: 这是一个断言（Assertion）语句。它检查传入的 stride 值是否属于列表 [1, 2] 中的一个。
        # 作用: 参数校验，确保传入的 stride 是有效的，防止后续计算出错。

        exp_channel = in_channel * expand_ratio
        self.use_res_connect = use_res_connect
        self.inv_res = Sequential(
            Conv_Block(in_channel=in_channel, out_channel=exp_channel, kernel_size=1, stride=1, padding=0),
            Conv_Block(in_channel=exp_channel, out_channel=exp_channel, kernel_size=3, stride=stride, padding=1,
                       group=exp_channel),
            Conv_Block(in_channel=exp_channel, out_channel=out_channel, kernel_size=1, stride=1, padding=0,
                       is_linear=True)
            # 在将高维特征压缩回低维（Bottleneck）时，如果使用 ReLU 这样的非线性激活函数，会破坏由于维度压缩而仅存的特征信息。因此，在最后一层去掉了激活函数，使用线性输出。
        )

    def forward(self, x):
        if self.use_res_connect:
            return x + self.inv_res(x)
        # 条件：在 MobileNetV2 的标准定义中，只有当 Stride=1 且 输入通道数等于输出通道数 时，才使用残差连接。这个连接帮助梯度在深层网络中更好地传播，避免梯度消失。
        else:
            return self.inv_res(x)

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

    def forward(self, attribute_gt, landmark_gt, euler_angle_gt, angle,
                landmarks, train_batchsize):
        '''
        forward 的 Docstring

        :param self: 说明
        :param attribute_gt: 类型 torch.Tensor, 形状 (batch_size, n_attributes)包含每个样本属性的张量，例如性别、年龄
        :param landmark_gt: 类型 torch.Tensor, 形状 (batch_size, n_landmarks * 2 )包含每个样本地标点的张量，例如(x1, y1, x2, y2, ..., xN, yN)
        :param euler_angle_gt: 类型 torch.Tensor, 形状 (batch_size, 3)包含每个样本欧拉角的张量（俯仰角、偏航）角、滚转角
        :param angle: 类型 torch.Tensor, 形状 (batch_size, 3)包含预测的欧拉角的张量
        :param landmarks: 类型 torch.Tensor, 形状 (batch_size, n_landmarks * 2)包含预测的地标点的张量
        :param train_batchsize: 类型 int, 训练时的批量大小
        '''
        weight_angle = torch.sum(1 - torch.cos(angle - euler_angle_gt), axis=1)
        # 计算几何信息：Σ(1-cosθn^k)  k=1,2,3
        # 最终得到每张图片的角度权重，形状为 (batch_size, 1)
        attributes_w_n = attribute_gt[:, 1:6].float()
        # 计算属性权重矩阵，其中第1行到第5行分别表示侧脸、正脸、抬头、低头、表情/遮挡等属性。
        # 其元素的数值为{0, 1}表示二分类，即“有无此属性”
        mat_ratio = torch.mean(attributes_w_n, axis=0)
        # 计算每个属性在当前批次中的平均值，亦或者者说频率
        # 因为每个元素的值只能是0（没有）或1（有），因此平均值就代表了当前属性在样本批次中出现的频率
        mat_ratio = torch.Tensor([
            1.0 / (x) if x > 0 else train_batchsize for x in mat_ratio
        ]).to(device)
        # 计算倒数权重，因为频率越低，要求惩罚的权重越大
        # 如果这样不处理，网络会倾向于只学习好占多数的简单样本，而忽略少数困难样本，导致 loss 被简单样本主导。
        weight_attribute = torch.sum(attributes_w_n.mul(mat_ratio), axis=1)
        # .mul 不是矩阵乘法，而是带有广播机制的按元素相乘
        # 例如，假设 attributes_w_n 的形状是 (batch_size, 5)，mat_ratio 的形状是 (1, 5)，则 mat_ratio 会被广播成 (batch_size, 5)，然后逐元素相乘。
        # 结果就是，在属性矩阵中，如果一个样本具有某个属性，则该属性的标称值从原来的1变为该属性的倒数权重，从而增加了该样本在总损失中的贡献。
        # 每个属性明码标价
        # 然后，将这些加权后的属性值相加，得到每个样本的总属性权重。最终形状： (batch_size, 1)
        # 这一行代码的总体作用就是，为当前 Batch 中的每一张图片，根据它包含的属性，累加计算出该图片的最终 Loss 权重，包含的困难属性越多，越稀有，这张图在计算Loss时所占的比重越大。 

        l2_distant = torch.sum(
            (landmark_gt - landmarks) * (landmark_gt - landmarks), axis=1)
        # 计算每个样本的地标点 L2 距离的平方和，形状为 (batch_size, 1)，每一行元素的形式为：
        # x1²+y1²+x2²+y2²+...+xN²+yN²
        return torch.mean(weight_angle * weight_attribute *
                          l2_distant), torch.mean(l2_distant)
        # 这里同样不是矩阵乘法，而是按元素相乘，最终第一项得到完整的 Loss 函数值
        # 最后求均值而不是求和，如果是求和的话，Loss 会随着 Batch Size 的增大而增大，导致超参数不稳定
        # 为什么还要返回未加权的 L2 距离平方和的均值呢？这是给人看的，用于监控指标（Metric / Monitoring）。它反映了模型当前预测的坐标和真实坐标平均相差多少。因为第一个 Loss 被权重“污染”了，你无法通过它判断模型到底收敛没有。
        # 也就是：如果只看第一个 Loss，你不知道 Loss 变大是因为模型变差

def smoothL1(y_true, y_pred, beta=1):
    """
    very similar to the smooth_l1_loss from pytorch, but with
    the extra beta parameter
    """
    mae = torch.abs(y_true - y_pred)
    loss = torch.sum(torch.where(mae > beta, mae - 0.5 * beta, 0.5 * mae**2 / beta), axis=-1)
    # 大误差用L1 Loss：mae - 0.5 * beta，这部分的梯度是常数（1 或 -1），防止梯度爆炸。当碰到离群点（Outliers）时，不会因为误差非常大而产生巨大的梯度把模型参数打乱
    # 小误差类似L2 Loss：0.5 * mae**2 / beta，这部分的梯度在原点附近是动态减小的（越来越接近 0），能够平滑趋近于零。L1 Loss 在 0 点不可导且梯度始终为 1，容易在最优解附近震荡无法收敛，Smooth L1 解决了这个问题。
    return torch.mean(loss)

def wing_loss(y_true, y_pred, w=10.0, epsilon=2.0, N_LANDMARK=106):
    y_pred = y_pred.reshape(-1, N_LANDMARK, 2)
    y_true = y_true.reshape(-1, N_LANDMARK, 2)
    # 将输入数据恢复成标准的集合形状
    # 神经网络的输出通常是摊平的一维向量，形状为 (batch_size, N_LANDMARK * 2)
    # 本操作将其重新调整为 (batch_size, N_LANDMARK, 2)

    x = y_true - y_pred
    c = w * (1.0 - math.log(1.0 + w / epsilon))
    absolute_x = torch.abs(x)
    # 连续性常数：C = w[1-ln(1+w/ε)]
    losses = torch.where(w > absolute_x,
                         w * torch.log(1.0 + absolute_x / epsilon),
                         absolute_x - c)
    # 小误差区间(w > |x|)：wln(1+|x|/ε)，类似于 L2 Loss，在原点附近平滑收敛
    # 大误差区间(|x| >= w)：|x| - C，类似于 L1 Loss，防止梯度爆炸（离群点）
    loss = torch.mean(torch.sum(losses, axis=[1, 2]), axis=0)
    # torch.sum(losses, axis=[1, 2]): 沿着第1维和第2维求和，即得到每一张图片的总误差
    # torch.mean(..., axis=0): 最后对所有图片的总误差求均值，得到最终的 Loss 值
    return loss

class GhostModule(Module):
    '''
    内部操作：标准卷积采用Pointwise Conv, 剩余的 feature maps 通过 Depthwise Conv 生成
    '''
    def __init__(self, in_channel, out_channel, is_linear=False):
        super(GhostModule, self).__init__()
        self.out_channel = out_channel
        init_channel = math.ceil(out_channel / 2)
        # 先利用标准卷积生成一半数量的 feature maps
        new_channel = init_channel
        # 需要利用用标准卷积生成的 feature maps 通过廉价操作（如深度卷积）生成剩余的 feature maps

        self.primary_conv = Conv_Block(in_channel, init_channel, 1, 1, 0, is_linear=is_linear)
        self.cheap_operation = Conv_Block(init_channel, new_channel, 3, 1, 1, group=init_channel, is_linear=is_linear)
        # group=init_channel 表示这是一个深度卷积（Depthwise Convolution），即廉价操作

    def forward(self, x):
        x1 = self.primary_conv(x)
        x2 = self.cheap_operation(x1)
        out = torch.cat([x1, x2], dim=1)
        return out[:, :self.out_channel, :, :]
        # 切片的意思是保留前面的 out_channel 个通道，丢弃多余的通道，因为在out_channel为奇数时会多生成一个通道
        # PyTorch 的 Tensor 维度： N, C, H, W

class GhostBottleneck(Module):
    def __init__(self, in_channel, hidden_channel, out_channel, stride):
        super(GhostBottleneck, self).__init__()
        assert stride in [1, 2]

        self.ghost_conv = Sequential(
            # GhostModule, 升维(Channels)
            GhostModule(in_channel, hidden_channel, is_linear=False),
            # DepthwiseConv-linear, 3×3 DWConv
            # stride = 1时，不做这一层 DWConv
            Conv_Block(hidden_channel, hidden_channel, 3, stride, 1, group=hidden_channel,
                       is_linear=True) if stride == 2 else Sequential(),
            # GhostModule-linear, 降维(Channels)，最后一层不做 ReLU 激活
            GhostModule(hidden_channel, out_channel, is_linear=True)
        )

        if stride == 1 and in_channel == out_channel:
            self.shortcut = Sequential()
        else:
            # shortcut 不可以直接相加时，使用 DWConv + PWConv 来调整维度和尺寸
            self.shortcut = Sequential(
                Conv_Block(in_channel, in_channel, 3, stride, 1, group=in_channel, is_linear=True),
                Conv_Block(in_channel, out_channel, 1, 1, 0, is_linear=True)
            )

    def forward(self, x):
        return self.ghost_conv(x) + self.shortcut(x)

class PFLD_GhostNet(Module):
    def __init__(self, width_factor=1, input_size=112, landmark_number=98):
        super(PFLD_GhostNet, self).__init__()

        self.conv1 = Conv_Block(3, int(64 * width_factor), 3, 2, 1)
        self.conv2 = Conv_Block(int(64 * width_factor), int(64 * width_factor), 3, 1, 1, group=int(64 * width_factor))

        self.conv3_1 = GhostBottleneck(int(64 * width_factor), int(128 * width_factor), int(80 * width_factor), stride=2)
        self.conv3_2 = GhostBottleneck(int(80 * width_factor), int(160 * width_factor), int(80 * width_factor), stride=1)
        self.conv3_3 = GhostBottleneck(int(80 * width_factor), int(160 * width_factor), int(80 * width_factor), stride=1)

        self.conv4_1 = GhostBottleneck(int(80 * width_factor), int(240 * width_factor), int(96 * width_factor), stride=2)
        self.conv4_2 = GhostBottleneck(int(96 * width_factor), int(288 * width_factor), int(96 * width_factor), stride=1)
        self.conv4_3 = GhostBottleneck(int(96 * width_factor), int(288 * width_factor), int(96 * width_factor), stride=1)

        self.conv5_1 = GhostBottleneck(int(96 * width_factor), int(384 * width_factor), int(144 * width_factor), stride=2)
        self.conv5_2 = GhostBottleneck(int(144 * width_factor), int(576 * width_factor), int(144 * width_factor), stride=1)
        self.conv5_3 = GhostBottleneck(int(144 * width_factor), int(576 * width_factor), int(144 * width_factor), stride=1)
        self.conv5_4 = GhostBottleneck(int(144 * width_factor), int(576 * width_factor), int(144 * width_factor), stride=1)

        self.conv6 = GhostBottleneck(int(144 * width_factor), int(288 * width_factor), int(16 * width_factor), stride=1)
        self.conv7 = Conv_Block(int(16 * width_factor), int(32 * width_factor), 3, 1, 1)
        self.conv8 = Conv_Block(int(32 * width_factor), int(128 * width_factor), input_size // 16, 1, 0, has_bn=False)

        self.avg_pool1 = AvgPool2d(input_size // 2)
        self.avg_pool2 = AvgPool2d(input_size // 4)
        self.avg_pool3 = AvgPool2d(input_size // 8)
        self.avg_pool4 = AvgPool2d(input_size // 16)

        self.fc = Linear(int(512 * width_factor), landmark_number * 2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x1 = self.avg_pool1(x)
        x1 = x1.view(x1.size(0), -1)
        # 保持第0个维度不变，其余展平
        # 即：将形状为 (Batch_Size, Channels, Height, Width) 的张量展平为 (Batch_Size, Channels × Height × Width)

        x = self.conv3_1(x)
        x = self.conv3_2(x)
        x = self.conv3_3(x)
        x2 = self.avg_pool2(x)
        x2 = x2.view(x2.size(0), -1)

        x = self.conv4_1(x)
        x = self.conv4_2(x)
        x = self.conv4_3(x)
        x3 = self.avg_pool3(x)
        x3 = x3.view(x3.size(0), -1)

        x = self.conv5_1(x)
        x = self.conv5_2(x)
        x = self.conv5_3(x)
        x = self.conv5_4(x)
        x4 = self.avg_pool4(x)
        x4 = x4.view(x4.size(0), -1)

        x = self.conv6(x)
        x = self.conv7(x)
        x5 = self.conv8(x)
        x5 = x5.view(x5.size(0), -1)

        multi_scale = torch.cat([x1, x2, x3, x4, x5], 1)
        landmarks = self.fc(multi_scale)

        return landmarks

class MobileOneBlock(nn.Module):
    def __init__(self, ...):
        # ...
        if inference_mode:
            # 推理模式：只定义一个卷积层
            self.reparam_conv = nn.Conv2d(...)
        else:
            # 训练模式：定义多分支
            # 1. Skip Branch (Identity + BN)
            self.rbr_skip = nn.BatchNorm2d(...) if ... else None
            # 2. Conv Branches (Over-parameterized)
            self.rbr_conv = nn.ModuleList([...])
            # 3. Scale Branch (1x1 Conv)
            self.rbr_scale = self._conv_bn(kernel_size=1, ...)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 推理模式：单层卷积
        if self.inference_mode:
            return self.activation(self.se(self.reparam_conv(x)))

        # 训练模式：多分支输出相加
        # Result = Skip(x) + Scale(x) + Conv_Branches(x)
        identity_out = self.rbr_skip(x) if self.rbr_skip is not None else 0
        scale_out = self.rbr_scale(x) if self.rbr_scale is not None else 0
        out = scale_out + identity_out
        for ix in range(self.num_conv_branches):
            out += self.rbr_conv[ix](x)
        return self.activation(self.se(out))

    def reparameterize(self):
        # 将多分支参数融合，切换为推理模式
        kernel, bias = self._get_kernel_bias()
        self.reparam_conv = nn.Conv2d(...)
        self.reparam_conv.weight.data = kernel
        self.reparam_conv.bias.data = bias
        self.inference_mode = True

class GhostOneModule(Module):
    def __init__(self, ...):
        # 1. Primary Conv (主特征): 使用 MobileOneBlock 替代标准 1x1 Conv
        self.primary_conv = MobileOneBlock(..., kernel_size=1)
        # 2. Cheap Operation (廉价特征生成): 使用 MobileOneBlock 替代标准 3x3 DWConv
        self.cheap_operation = MobileOneBlock(..., kernel_size=3, groups=half_outchannel)

    def forward(self, x):
        x1 = self.primary_conv(x)
        x2 = self.cheap_operation(x1)
        # 拼接主特征和生成特征
        out = torch.cat([x1, x2], dim=1)
        return out

class GhostOneBottleneck(Module):
    def __init__(self, ...):
        # 移除了 Residual Connection (Shortcut)
        # 仅包含堆叠的 GhostOneModule 和下采样层
        self.ghost_conv = Sequential(
            # 1. 升维 (Expansion)
            GhostOneModule(in_channel, hidden_channel, ...),

            # 2. 深度卷积下采样 (仅当 stride=2 时)
            MobileOneBlock(..., stride=stride, ...) if stride == 2 else Sequential(),

            # 3. 降维 (Projection)
            GhostOneModule(hidden_channel, out_channel, is_linear=True, ...)
        )

    def forward(self, x):
        # 直接输出，没有 + x 操作
        return self.ghost_conv(x)

def _fuse_bn_tensor(self, branch) -> Tuple[torch.Tensor, torch.Tensor]:
    """ 分支融合原子操作：将 (Conv+BN) 或 (BN) 融合为 (Conv_weight, Conv_bias)。

    原理：
    BN 公式: y = (x - mean) / sqrt(var + eps) * gamma + beta
    卷积公式: y = Wx 
    融合后: y = (W * gamma / std) * x + (beta - mean * gamma / std)

    :param branch: 输入分支，可能是 nn.Sequential(Conv, BN) 或者单独的 nn.BatchNorm2d
    """
    if isinstance(branch, nn.Sequential):
        # Case 1: 分支是 Conv + BN
        # 为什么是 .conv 和 .bn？详见本类的辅助函数 .conv_bn 的命名约定
        kernel = branch.conv.weight
        running_mean = branch.bn.running_mean
        running_var = branch.bn.running_var
        gamma = branch.bn.weight
        beta = branch.bn.bias
        eps = branch.bn.eps
    else:
        # Case 2: 分支只有 BN (Skip Connection)
        assert isinstance(branch, nn.BatchNorm2d)
        if not hasattr(self, 'id_tensor'):
            # 构造一个恒等映射卷积核（Identity Kernel）
            # 这是一个 KxK 的卷积核，除了中心点是 1，其余都是 0
            input_dim = self.in_channels // self.groups
            # 这是考虑了 分组卷积 (Group Convolution) 的情况。
            # 如果是标准卷积，则self.groups = 1, 每个卷积核的深度等于输入通道数self.in_channels；如果是DWConv，则self.groups = self.in_channels, 每个卷积核的深度为1。
            kernel_value = torch.zeros((self.in_channels,
                                        # 按道理这里应该是self.out_channels，但是对应 Case 2: 分支只有 BN (Skip Connection)，这是恒等映射，self.in_channels == self.out_channels
                                        input_dim,
                                        self.kernel_size,
                                        self.kernel_size),
                                        dtype=branch.weight.dtype,
                                        device=branch.weight.device)
            # 构建了一个标准的 PyTorch 卷积权重容器，形状为：()[out_channels, in_channels/groups, K, K]，初始值全为0
            # 如果是标准卷积，则 input_dim = self.in_channels；对应标准卷积的卷积核需要处理所有通道；如果是 DWConv，则 input_dim = 1；对应 DWConv 的卷积核每个只处理一个通道。
            for i in range(self.in_channels):
                # 按道理这里应该是self.out_channels，但是对应 Case 2: 分支只有 BN (Skip Connection)，这是恒等映射，self.in_channels == self.out_channels
                kernel_value[i, i % input_dim,
                                self.kernel_size // 2,
                                self.kernel_size // 2] = 1
                # 让卷积核中间的值为0
                # 如果是标准卷积，input_dim = self.in_channels, 则 i % input_dim = i，即每隔通道的[self.kernel_size // 2, self.kernel_size // 2]位置为1
                # 如果是 DWConv，input_dim = 1, 则 i % input_dim = 0，即只有每个卷积核的第0个输入通道的[self.kernel_size // 2, self.kernel_size // 2]位置为1
            self.id_tensor = kernel_value
        kernel = self.id_tensor
        running_mean = branch.running_mean
        running_var = branch.running_var
        gamma = branch.weight
        beta = branch.bias
        eps = branch.eps

    # 融合公式实现
    std = (running_var + eps).sqrt()
    # t = gamma / std
    t = (gamma / std).reshape(-1, 1, 1, 1) # reshape 以支持广播乘法

    # 新 Weight = 旧 Weight * (gamma / std)
    # 新 Bias = beta - mean * (gamma / std)
    return kernel * t, beta - running_mean * gamma / std