You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
97 lines
3.2 KiB
97 lines
3.2 KiB
# Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserve. |
|
# |
|
# Licensed under the Apache License, Version 2.0 (the "License"); |
|
# you may not use this file except in compliance with the License. |
|
# You may obtain a copy of the License at |
|
# |
|
# http://www.apache.org/licenses/LICENSE-2.0 |
|
# |
|
# Unless required by applicable law or agreed to in writing, software |
|
# distributed under the License is distributed on an "AS IS" BASIS, |
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|
# See the License for the specific language governing permissions and |
|
# limitations under the License. |
|
|
|
import paddle |
|
import paddle.nn as nn |
|
import paddle.nn.functional as F |
|
|
|
|
|
def pdist(e, squared=False, eps=1e-12): |
|
e_square = e.pow(2).sum(axis=1) |
|
prod = paddle.mm(e, e.t()) |
|
res = (e_square.unsqueeze(1) + e_square.unsqueeze(0) - 2 * prod).clip( |
|
min=eps) |
|
|
|
if not squared: |
|
res = res.sqrt() |
|
return res |
|
|
|
|
|
class RKdAngle(nn.Layer): |
|
# reference: https://github.com/lenscloth/RKD/blob/master/metric/loss.py |
|
def __init__(self, target_size=None): |
|
super().__init__() |
|
if target_size is not None: |
|
self.avgpool = paddle.nn.AdaptiveAvgPool2D(target_size) |
|
else: |
|
self.avgpool = None |
|
|
|
def forward(self, student, teacher): |
|
# GAP to reduce memory |
|
if self.avgpool is not None: |
|
# NxC1xH1xW1 -> NxC1x1x1 |
|
student = self.avgpool(student) |
|
# NxC2xH2xW2 -> NxC2x1x1 |
|
teacher = self.avgpool(teacher) |
|
|
|
# reshape for feature map distillation |
|
bs = student.shape[0] |
|
student = student.reshape([bs, -1]) |
|
teacher = teacher.reshape([bs, -1]) |
|
|
|
td = (teacher.unsqueeze(0) - teacher.unsqueeze(1)) |
|
norm_td = F.normalize(td, p=2, axis=2) |
|
t_angle = paddle.bmm(norm_td, norm_td.transpose([0, 2, 1])).reshape( |
|
[-1, 1]) |
|
|
|
sd = (student.unsqueeze(0) - student.unsqueeze(1)) |
|
norm_sd = F.normalize(sd, p=2, axis=2) |
|
s_angle = paddle.bmm(norm_sd, norm_sd.transpose([0, 2, 1])).reshape( |
|
[-1, 1]) |
|
loss = F.smooth_l1_loss(s_angle, t_angle, reduction='mean') |
|
return loss |
|
|
|
|
|
class RkdDistance(nn.Layer): |
|
# reference: https://github.com/lenscloth/RKD/blob/master/metric/loss.py |
|
def __init__(self, eps=1e-12, target_size=1): |
|
super().__init__() |
|
self.eps = eps |
|
if target_size is not None: |
|
self.avgpool = paddle.nn.AdaptiveAvgPool2D(target_size) |
|
else: |
|
self.avgpool = None |
|
|
|
def forward(self, student, teacher): |
|
# GAP to reduce memory |
|
if self.avgpool is not None: |
|
# NxC1xH1xW1 -> NxC1x1x1 |
|
student = self.avgpool(student) |
|
# NxC2xH2xW2 -> NxC2x1x1 |
|
teacher = self.avgpool(teacher) |
|
|
|
bs = student.shape[0] |
|
student = student.reshape([bs, -1]) |
|
teacher = teacher.reshape([bs, -1]) |
|
|
|
t_d = pdist(teacher, squared=False) |
|
mean_td = t_d.mean() |
|
t_d = t_d / (mean_td + self.eps) |
|
|
|
d = pdist(student, squared=False) |
|
mean_d = d.mean() |
|
d = d / (mean_d + self.eps) |
|
|
|
loss = F.smooth_l1_loss(d, t_d, reduction="mean") |
|
return loss
|
|
|