面向目标姿态估计的模板匹配模型训练异常问询

我正在测试一款用于匹配已知目标查询图像与对应模板图像的模型，目标姿态可能一致（需处理对称目标与严重遮挡场景，因此常为一对多关系）。
模型输入为查询图像+候选模板图像的图像对，输出0.0表示姿态不一致，1.0表示姿态一致，采用L1_loss训练。
训练使用合成数据批量输入，每个查询图像对应：

• 正例：查询图像+关联的正确模板图像（期望输出1.0）
  [正例图像示例]
• 负例：查询图像+随机模板图像（期望输出0.0）
  [负例图像示例]

异常点在于：当负样本为正模板的平面旋转图像时，模型训练收敛快，表现极佳（正例平均分类值≈0.99，负例≈0.1）；但当负样本为任意3D姿态的随机模板时，模型表现极差（正例平均分类值≈0.75，负例≈0.5），而理论上此类正负例差异更大，应更易区分。

模型 #

class TemplateEvaluator(nn.Module):
    def __init__(self, q_encoder=resnet18(weights=ResNet18_Weights.IMAGENET1K_V1), t_encoder=resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)):
        super(TemplateEvaluator, self).__init__()
        self.q_encoder = q_encoder
        self.t_encoder = t_encoder
        
        self.fc = nn.Sequential(
            nn.Linear(2000, 1),
            nn.Sigmoid()
        )
    
    def forward(self, data):
        q = data[0]
        t = data[1]
        q = self.q_encoder(q)
        t = self.t_encoder(t)
        res = self.fc(torch.cat([q,t],-1))
        return res

训练步骤 #

• cb_id包含关联的正确模板ID（角度差最小的模板）
• t_img_rand为负样本模板

def template_eval_train_step(iteration, models, data, codebook, opts=None, show=False, metric_label=''):
    # Get query image, associated codebook template ID, and associated orientation
    q_img, cb_id, rot = data
    n = q_img.shape[0]
    t_eval = models[0]
    
    # Get random template IDs
    cb_id_rand = np.random.choice(codebook["size"],n)

    # Get associated and random template images
    t_img = torch.stack([cb_get_img(i,codebook) for i in cb_id]).to(device)

    # Uncomment to use random template as neg cases
    t_img_rand = torch.stack([cb_get_img(i,codebook) for i in cb_id_rand]).to(device)
    # Uncomment to use in-plane rotations of pos template as neg cases
    # t_img_rand = torch.stack([rotate_image_tensor(y,np.random.random()*360) for y in t_img])
    
    # Cases with similar template image ('Positive')
    p_cases = torch.stack([q_img.permute(0, 3, 1, 2),t_img.permute(0, 3, 1, 2)])

    # Cases with random template image ('Negative')
    n_cases = torch.stack([q_img.permute(0, 3, 1, 2),t_img_rand.permute(0, 3, 1, 2)])

    # Mix together for 50/50 distribution in batch
    mixed_cases = torch.concat([p_cases,n_cases], 1)

    # Run model
    c = t_eval(mixed_cases)

    # Get classification for pos and neg cases
    p_cls = c[:n]
    n_cls = c[n:]

    # Compute loss
    p_loss = F.l1_loss(p_cls, torch.ones_like(p_cls, requires_grad=True))
    n_loss = F.l1_loss(n_cls, torch.zeros_like(n_cls, requires_grad=True))
    loss = (p_loss + n_loss)/2
    
    # Visualise pos and neg case at i=0
    if show:
        i=0
        view([q_img[i].detach().cpu().numpy(), t_img[i].detach().cpu().numpy()])
        print("p_cls:",p_cls[i].detach().cpu().numpy())
        view([q_img[i].detach().cpu().numpy(), t_img_rand[i].detach().cpu().numpy()])
        print("n_cls:",n_cls[i].detach().cpu().numpy())

    # Run optimizer (if given)
    if opts is not None:
        opts[0].zero_grad()
        loss.backward()
        
        # Print gradient info
        if show:
            t_eval.cpu()
            plot_grad_flow(t_eval.named_parameters())
            t_eval.to(device)
        
        opts[0].step()

    # Compute eval metrics
    p_rate = p_cls.sum() / n
    n_rate = n_cls.sum() / n

    # Garbage collection 
    gc.collect()

    return [ {"label": metric_label, "name": "loss", "value":loss.cpu().item()},
             {"label": metric_label, "name": "p_rate", "value":p_rate.cpu().item()},
             {"label": metric_label, "name": "n_rate", "value":n_rate.cpu().item()}]

训练循环 #

• init_train、init_verify用于切换模型训练/评估模式
• train_step为上述训练步骤函数

def fit(epochs, models, init_train, init_verify, train_step, verify_step, opts, train_dl, verify_dl, eval_dl, codebook, vis_epoch_step=10):
    train_data = []
    verify_data = []
    eval_data = []

    for epoch in tqdm(range(epochs)):
        init_train(epoch, models)
        
        i = 0
        for data in train_dl:
            train_metrics = train_step(epoch, models, data, opts=opts, codebook=codebook, show=epoch % vis_epoch_step == 0 and i == 0)
            train_data.append(train_metrics)
            i = i + 1
            
            n = len(train_dl)
            p = round((i/n)*100)
            if p>0:
                sys.stdout.write('\r')
                bar_len = round(p/5)
                empty_len = round((100-p)/5)
                sys.stdout.write("Train batch %d/%d [%s%s] %d%%" % (i, n, '#'*bar_len, '_'*empty_len, p))
                sys.stdout.flush()
            
        # verification step
        init_verify(epoch, models)
        with torch.no_grad():
            
            i = 0
            for data in verify_dl:
                verify_metrics = verify_step(epoch, models, data, codebook=codebook, show=epoch % vis_epoch_step == 0 and i == 0)
                verify_data.append(verify_metrics)
                i = i + 1
            
                n = len(verify_dl)
                p = round((i/n)*100)
                if p>0:
                    sys.stdout.write('\r')
                    bar_len = round(p/5)
                    empty_len = round((100-p)/5)
                    sys.stdout.write("Verification batch %d/%d [%s%s] %d%%" % (i, n, '#'*bar_len, '_'*empty_len, p))
                    sys.stdout.flush()
...

正模板平面旋转为负样本的情况 #

使用以下代码生成负样本：

t_img_rand = torch.stack([rotate_image_tensor(y,np.random.random()*360) for y in t_img])

训练结果（p/n_rate为正负例平均分类值）：
[训练结果曲线图像]

示例：
[正例图像示例]
p_cls: [0.998]

[负例图像示例（平面旋转模板）]
n_cls: [0.000]

随机模板为负样本的情况 #

使用以下代码生成负样本：

cb_id_rand = np.random.choice(codebook["size"],n)
t_img_rand = torch.stack([cb_get_img(i,codebook) for i in cb_id_rand]).to(device)

训练结果（p/n_rate为正负例平均分类值）：
[训练结果曲线图像]

示例：
[正例图像示例]
p_cls: [0.001]

[负例图像示例（随机3D姿态模板）]
n_cls: [0.998]

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