|
|
|
|
| from __future__ import print_function |
| import torch |
| import torch.nn as nn |
| from .model import* |
| import torch.nn.functional as F |
|
|
| def vq_cos_sim(embedding, input_tensor, use_dynamic_p=False,ddp=False): |
|
|
| if ddp: |
| embedding_weight = embedding.module.weight |
| else: |
| embedding_weight = embedding.weight |
|
|
| |
| input_norm = F.normalize(input_tensor, p=2, dim=2) |
| embedding_norm = F.normalize(embedding_weight, p=2, dim=1) |
|
|
| similarity = torch.matmul(input_norm, embedding_norm.t()) |
| cos_sim_values, indices = similarity.max(dim=2) |
|
|
| if use_dynamic_p: |
| |
| return indices.squeeze(), cos_sim_values.squeeze() |
| |
| return indices.squeeze() |
|
|
|
|
| class RatioLossWithMSELoss(nn.Module): |
| def __init__(self, total_iters, min_weight=0.001, max_weight=1,eps=torch.tensor(1e-3, dtype=torch.bfloat16)): |
| super(RatioLossWithMSELoss, self).__init__() |
| self.eps = eps |
| self.total_iters = total_iters |
| self.min_weight = min_weight |
| self.max_weight = max_weight |
| self.mse=nn.MSELoss() |
| self.losses={} |
| def forward(self, output, target, current_iter): |
|
|
| weight = self.min_weight + (self.max_weight - self.min_weight) * (current_iter / self.total_iters) |
| loss = (torch.abs(target - output)) / (torch.abs(target) + self.eps) |
| weighted_loss = weight * loss |
|
|
| self.losses['ratio']=loss.mean() |
| self.losses['mse']=self.mse(output,target) |
| return weighted_loss.mean()+self.mse(output,target) |
| |
|
|
|
|
|
|
| if __name__=='__main__': |
| pass |
