| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768 |
- import torch
- import torch.nn.functional as F
- from torch import nn
- class CELoss(nn.Module):
- def __init__(self,
- smoothing=False,
- with_all=False,
- ignore_index=-1,
- **kwargs):
- super(CELoss, self).__init__()
- if ignore_index >= 0:
- self.loss_func = nn.CrossEntropyLoss(reduction='mean',
- ignore_index=ignore_index)
- else:
- self.loss_func = nn.CrossEntropyLoss(reduction='mean')
- self.smoothing = smoothing
- self.with_all = with_all
- def forward(self, pred, batch):
- pred = pred['res']
- if isinstance(pred, dict): # for ABINet
- loss = {}
- loss_sum = []
- for name, logits in pred.items():
- if isinstance(logits, list):
- logit_num = len(logits)
- all_tgt = torch.cat([batch[1]] * logit_num, 0)
- all_logits = torch.cat(logits, 0)
- flt_logtis = all_logits.reshape([-1, all_logits.shape[2]])
- flt_tgt = all_tgt.reshape([-1])
- else:
- flt_logtis = logits.reshape([-1, logits.shape[2]])
- flt_tgt = batch[1].reshape([-1])
- loss[name + '_loss'] = self.loss_func(flt_logtis, flt_tgt)
- loss_sum.append(loss[name + '_loss'])
- loss['loss'] = sum(loss_sum)
- return loss
- else:
- if self.with_all: # for ViTSTR
- tgt = batch[1]
- pred = pred.reshape([-1, pred.shape[2]])
- tgt = tgt.reshape([-1])
- loss = self.loss_func(pred, tgt)
- return {'loss': loss}
- else: # for NRTR
- max_len = batch[2].max()
- tgt = batch[1][:, 1:2 + max_len]
- pred = pred.reshape([-1, pred.shape[2]])
- tgt = tgt.reshape([-1])
- if self.smoothing:
- eps = 0.1
- pred.shape[1]
- one_hot = F.one_hot(tgt, pred.shape[1])
- one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (-1)
- log_prb = F.log_softmax(pred, dim=1)
- non_pad_mask = torch.not_equal(
- tgt,
- torch.zeros(tgt.shape,
- dtype=tgt.dtype,
- device=tgt.device))
- loss = -(one_hot * log_prb).sum(dim=1)
- loss = loss.masked_select(non_pad_mask).mean()
- else:
- loss = self.loss_func(pred, tgt)
- return {'loss': loss}
|
