重训练FaceNet：如何冻结预训练网络层并替换自定义损失函数

已完成基础代码 #

import torch
from facenet_pytorch import InceptionResnetV1, MTCNN
from torch.utils.data import DataLoader
from torchvision import datasets
import numpy as np
import pandas as pd
import os

workers = 0 if os.name == 'nt' else 4
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('Running on device: {}'.format(device))

mtcnn = MTCNN(
    image_size=160, margin=0, min_face_size=20,
    thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True,
    device=device
)

核心问题实现方法 #

1. 指定层冻结（固定特征提取层参数） #

层冻结的核心逻辑是关闭不需要调优层的梯度更新权限，仅让待微调层参与反向传播参数更新，操作分两步：

• 第一步：关闭全局梯度后按需解冻目标层
  先把MTCNN所有参数的requires_grad设为False，再找到需要微调的最后几层，重新打开梯度。操作前先执行print(mtcnn)打印完整网络结构，确认要保留的特征提取层、要微调的输出层对应的模块名，避免错冻漏冻。
  参考代码：

  # 冻结全部参数
  for param in mtcnn.parameters():
      param.requires_grad = False
  
  # 示例：解冻O-Net的最后3个输出层（可根据自己的微调需求替换层名）
  trainable_layers = [mtcnn.onet.conv6_1, mtcnn.onet.conv6_2, mtcnn.onet.conv6_3]
  for layer in trainable_layers:
      for param in layer.parameters():
          param.requires_grad = True
  

• 第二步：优化器仅传入可训练参数
  初始化优化器时过滤掉requires_grad=False的参数，从根源上避免冻结层被更新：

  optimizer = torch.optim.Adam(
      filter(lambda p: p.requires_grad, mtcnn.parameters()),
      lr=1e-4
  )
  

注意：微调开始前必须调用mtcnn.train()将网络切换到训练模式，否则BatchNorm、Dropout等层会保持推理状态，微调结果会异常。

2. 自定义损失函数接入 #

MTCNN原生是三任务联合损失（人脸分类、边界框回归、关键点定位），替换自定义损失不需要修改网络源码，只需要在训练循环中替换损失计算逻辑即可，操作分两步：

• 第一步：定义自定义损失类
  根据自己的任务需求实现损失计算逻辑，不需要受原生损失的权重、计算规则限制：

  import torch.nn as nn
  import torch.nn.functional as F
  
  class CustomCriterion(nn.Module):
      def __init__(self, cls_weight=1.0, box_weight=0.5, landmark_weight=0.5):
          super().__init__()
          self.cls_weight = cls_weight
          self.box_weight = box_weight
          self.landmark_weight = landmark_weight
  
      def forward(self, cls_pred, box_pred, landmark_pred, cls_gt, box_gt, landmark_gt):
          # 此处可替换为任意需要的损失计算逻辑，比如Focal Loss、ArcFace Loss等
          cls_loss = F.binary_cross_entropy_with_logits(cls_pred, cls_gt)
          box_loss = F.smooth_l1_loss(box_pred, box_gt)
          landmark_loss = F.smooth_l1_loss(landmark_pred, landmark_gt)
          return self.cls_weight * cls_loss + self.box_weight * box_loss + self.landmark_weight * landmark_loss
  
  criterion = CustomCriterion()
  

• 第二步：训练循环中接入自定义损失
  训练时单独调用对应子网的前向逻辑拿到预测值，传入损失函数计算即可，不需要依赖MTCNN原生的检测前向流程：

  mtcnn.train()
  for epoch in range(epochs):
      for batch_imgs, batch_labels in train_loader:
          batch_imgs = batch_imgs.to(device)
          # 示例：仅微调O-Net时，直接调用O-Net前向得到三类预测输出
          cls_pred, box_pred, landmark_pred = mtcnn.onet(batch_imgs)
          # 用自定义损失计算总loss
          loss = criterion(
              cls_pred, box_pred, landmark_pred,
              batch_labels['cls'].to(device),
              batch_labels['box'].to(device),
              batch_labels['landmark'].to(device)
          )
          # 反向传播更新参数
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  

如果你的场景是冻结MTCNN做特征提取，只训练下游新增的分类/回归头，直接把MTCNN输出的特征传入下游头，损失对接下游头的输出和标签即可，逻辑更简单。

内容的提问来源于stack exchange，提问作者Norhther