| 34 | |
| 35 | |
| 36 | def consistency_loss(logits_s, logits_w, class_acc, it, ds, name='ce', T=1.0, p_cutoff=0.0, use_flex=False): |
| 37 | logits_w = logits_w.detach() |
| 38 | |
| 39 | if name == 'ce': |
| 40 | pseudo_label = torch.softmax(logits_w, dim=-1) |
| 41 | max_probs, max_idx = torch.max(pseudo_label, dim=-1) |
| 42 | if use_flex: |
| 43 | mask = max_probs.ge(p_cutoff * (class_acc[max_idx] / (2. - class_acc[max_idx]))).float() |
| 44 | else: |
| 45 | mask = max_probs.ge(p_cutoff).float() |
| 46 | select = max_probs.ge(p_cutoff).long() |
| 47 | |
| 48 | pseudo_label = torch.softmax(logits_w / T, dim=-1) |
| 49 | masked_loss = ce_loss(logits_s, pseudo_label, use_hard_labels=False) * mask |
| 50 | return masked_loss.mean(), mask.mean(), select, max_idx.long() |
| 51 | |
| 52 | if name == 'kld_tf': |
| 53 | # The implementation of loss_unsup for Google's TF |
| 54 | # logits_tgt by sharpening |
| 55 | logits_tgt = (logits_w / T).detach() |
| 56 | |
| 57 | # pseudo_labels_mask by confidence masking |
| 58 | pseudo_labels = torch.softmax(logits_w, dim=-1).detach() |
| 59 | p_class_max = torch.max(pseudo_labels, dim=-1, keepdim=False)[0] |
| 60 | loss_mask = p_class_max.ge(p_cutoff).float().detach() |
| 61 | |
| 62 | kld = F.kl_div(torch.log_softmax(logits_s, dim=-1), torch.softmax(logits_tgt, dim=-1), reduction='none') |
| 63 | masked_loss = kld * loss_mask.unsqueeze(dim=-1).repeat(1, pseudo_labels.shape[1]) |
| 64 | masked_loss = torch.sum(masked_loss, dim=1) |
| 65 | return masked_loss.mean(), loss_mask.mean() |
| 66 | |
| 67 | |
| 68 | else: |
| 69 | assert Exception('Not Implemented consistency_loss') |
| 70 | |
| 71 | |
| 72 | def torch_device_one(): |