| import torch |
| from torch import nn |
| from configs.paths_config import model_paths |
| from models.encoders.model_irse import Backbone |
|
|
|
|
| class IDLoss(nn.Module): |
| def __init__(self): |
| super(IDLoss, self).__init__() |
| print('Loading ResNet ArcFace') |
| self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se') |
| self.facenet.load_state_dict(torch.load(model_paths['ir_se50'])) |
| self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112)) |
| self.facenet.eval() |
|
|
| def extract_feats(self, x): |
| x = x[:, :, 35:223, 32:220] |
| x = self.face_pool(x) |
| x_feats = self.facenet(x) |
| return x_feats |
|
|
| def forward(self, y_hat, y, x, label=None, weights=None): |
| n_samples = x.shape[0] |
| x_feats = self.extract_feats(x) |
| y_feats = self.extract_feats(y) |
| y_hat_feats = self.extract_feats(y_hat) |
| y_feats = y_feats.detach() |
| total_loss = 0 |
| sim_improvement = 0 |
| id_logs = [] |
| count = 0 |
| for i in range(n_samples): |
| diff_target = y_hat_feats[i].dot(y_feats[i]) |
| diff_input = y_hat_feats[i].dot(x_feats[i]) |
| diff_views = y_feats[i].dot(x_feats[i]) |
|
|
| if label is None: |
| id_logs.append({'diff_target': float(diff_target), |
| 'diff_input': float(diff_input), |
| 'diff_views': float(diff_views)}) |
| else: |
| id_logs.append({f'diff_target_{label}': float(diff_target), |
| f'diff_input_{label}': float(diff_input), |
| f'diff_views_{label}': float(diff_views)}) |
|
|
| loss = 1 - diff_target |
| if weights is not None: |
| loss = weights[i] * loss |
|
|
| total_loss += loss |
| id_diff = float(diff_target) - float(diff_views) |
| sim_improvement += id_diff |
| count += 1 |
|
|
| return total_loss / count, sim_improvement / count, id_logs |
|
|
