Spaces:
Build error
Build error
File size: 2,592 Bytes
cb433d6 |
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 65 66 67 68 69 70 71 72 73 |
from fastai.vision import *
from modules.model import Model
class MultiLosses(nn.Module):
def __init__(self, one_hot=True):
super().__init__()
self.ce = SoftCrossEntropyLoss() if one_hot else torch.nn.CrossEntropyLoss()
self.bce = torch.nn.BCELoss()
@property
def last_losses(self):
return self.losses
def _flatten(self, sources, lengths):
return torch.cat([t[:l] for t, l in zip(sources, lengths)])
def _merge_list(self, all_res):
if not isinstance(all_res, (list, tuple)):
return all_res
def merge(items):
if isinstance(items[0], torch.Tensor): return torch.cat(items, dim=0)
else: return items[0]
res = dict()
for key in all_res[0].keys():
items = [r[key] for r in all_res]
res[key] = merge(items)
return res
def _ce_loss(self, output, gt_labels, gt_lengths, idx=None, record=True):
loss_name = output.get('name')
pt_logits, weight = output['logits'], output['loss_weight']
assert pt_logits.shape[0] % gt_labels.shape[0] == 0
iter_size = pt_logits.shape[0] // gt_labels.shape[0]
if iter_size > 1:
gt_labels = gt_labels.repeat(3, 1, 1)
gt_lengths = gt_lengths.repeat(3)
flat_gt_labels = self._flatten(gt_labels, gt_lengths)
flat_pt_logits = self._flatten(pt_logits, gt_lengths)
nll = output.get('nll')
if nll is not None:
loss = self.ce(flat_pt_logits, flat_gt_labels, softmax=False) * weight
else:
loss = self.ce(flat_pt_logits, flat_gt_labels) * weight
if record and loss_name is not None: self.losses[f'{loss_name}_loss'] = loss
return loss
def forward(self, outputs, *args):
self.losses = {}
if isinstance(outputs, (tuple, list)):
outputs = [self._merge_list(o) for o in outputs]
return sum([self._ce_loss(o, *args) for o in outputs if o['loss_weight'] > 0.])
else:
return self._ce_loss(outputs, *args, record=False)
class SoftCrossEntropyLoss(nn.Module):
def __init__(self, reduction="mean"):
super().__init__()
self.reduction = reduction
def forward(self, input, target, softmax=True):
if softmax: log_prob = F.log_softmax(input, dim=-1)
else: log_prob = torch.log(input)
loss = -(target * log_prob).sum(dim=-1)
if self.reduction == "mean": return loss.mean()
elif self.reduction == "sum": return loss.sum()
else: return loss
|