Spaces:
Running
on
Zero
Running
on
Zero
File size: 1,710 Bytes
a249588 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
# Copyright (c) OpenMMLab. All rights reserved.
import torch.nn as nn
from mmpose.registry import MODELS
@MODELS.register_module()
class FeaLoss(nn.Module):
"""PyTorch version of feature-based distillation from DWPose Modified from
the official implementation.
<https://github.com/IDEA-Research/DWPose>
Args:
student_channels(int): Number of channels in the student's feature map.
teacher_channels(int): Number of channels in the teacher's feature map.
alpha_fea (float, optional): Weight of dis_loss. Defaults to 0.00007
"""
def __init__(
self,
name,
use_this,
student_channels,
teacher_channels,
alpha_fea=0.00007,
):
super(FeaLoss, self).__init__()
self.alpha_fea = alpha_fea
if teacher_channels != student_channels:
self.align = nn.Conv2d(
student_channels,
teacher_channels,
kernel_size=1,
stride=1,
padding=0)
else:
self.align = None
def forward(self, preds_S, preds_T):
"""Forward function.
Args:
preds_S(Tensor): Bs*C*H*W, student's feature map
preds_T(Tensor): Bs*C*H*W, teacher's feature map
"""
if self.align is not None:
outs = self.align(preds_S)
else:
outs = preds_S
loss = self.get_dis_loss(outs, preds_T)
return loss
def get_dis_loss(self, preds_S, preds_T):
loss_mse = nn.MSELoss(reduction='sum')
N, C, H, W = preds_T.shape
dis_loss = loss_mse(preds_S, preds_T) / N * self.alpha_fea
return dis_loss
|