87 lines
2.6 KiB
Python
87 lines
2.6 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from torch.nn import functional as F
|
|
|
|
from dassl.optim import build_optimizer, build_lr_scheduler
|
|
from dassl.utils import count_num_param
|
|
from dassl.engine import TRAINER_REGISTRY, TrainerXU
|
|
from dassl.metrics import compute_accuracy
|
|
from dassl.modeling.ops import ReverseGrad
|
|
from dassl.engine.trainer import SimpleNet
|
|
|
|
|
|
class Prototypes(nn.Module):
|
|
|
|
def __init__(self, fdim, num_classes, temp=0.05):
|
|
super().__init__()
|
|
self.prototypes = nn.Linear(fdim, num_classes, bias=False)
|
|
self.temp = temp
|
|
|
|
def forward(self, x):
|
|
x = F.normalize(x, p=2, dim=1)
|
|
out = self.prototypes(x)
|
|
out = out / self.temp
|
|
return out
|
|
|
|
|
|
@TRAINER_REGISTRY.register()
|
|
class MME(TrainerXU):
|
|
"""Minimax Entropy.
|
|
|
|
https://arxiv.org/abs/1904.06487.
|
|
"""
|
|
|
|
def __init__(self, cfg):
|
|
super().__init__(cfg)
|
|
self.lmda = cfg.TRAINER.MME.LMDA
|
|
|
|
def build_model(self):
|
|
cfg = self.cfg
|
|
|
|
print("Building F")
|
|
self.F = SimpleNet(cfg, cfg.MODEL, 0)
|
|
self.F.to(self.device)
|
|
print("# params: {:,}".format(count_num_param(self.F)))
|
|
self.optim_F = build_optimizer(self.F, cfg.OPTIM)
|
|
self.sched_F = build_lr_scheduler(self.optim_F, cfg.OPTIM)
|
|
self.register_model("F", self.F, self.optim_F, self.sched_F)
|
|
|
|
print("Building C")
|
|
self.C = Prototypes(self.F.fdim, self.num_classes)
|
|
self.C.to(self.device)
|
|
print("# params: {:,}".format(count_num_param(self.C)))
|
|
self.optim_C = build_optimizer(self.C, cfg.OPTIM)
|
|
self.sched_C = build_lr_scheduler(self.optim_C, cfg.OPTIM)
|
|
self.register_model("C", self.C, self.optim_C, self.sched_C)
|
|
|
|
self.revgrad = ReverseGrad()
|
|
|
|
def forward_backward(self, batch_x, batch_u):
|
|
input_x, label_x, input_u = self.parse_batch_train(batch_x, batch_u)
|
|
|
|
feat_x = self.F(input_x)
|
|
logit_x = self.C(feat_x)
|
|
loss_x = F.cross_entropy(logit_x, label_x)
|
|
self.model_backward_and_update(loss_x)
|
|
|
|
feat_u = self.F(input_u)
|
|
feat_u = self.revgrad(feat_u)
|
|
logit_u = self.C(feat_u)
|
|
prob_u = F.softmax(logit_u, 1)
|
|
loss_u = -(-prob_u * torch.log(prob_u + 1e-5)).sum(1).mean()
|
|
self.model_backward_and_update(loss_u * self.lmda)
|
|
|
|
loss_summary = {
|
|
"loss_x": loss_x.item(),
|
|
"acc_x": compute_accuracy(logit_x, label_x)[0].item(),
|
|
"loss_u": loss_u.item(),
|
|
}
|
|
|
|
if (self.batch_idx + 1) == self.num_batches:
|
|
self.update_lr()
|
|
|
|
return loss_summary
|
|
|
|
def model_inference(self, input):
|
|
return self.C(self.F(input))
|