MCD.py

'''
Paper: Saito, K., Watanabe, K., Ushiku, Y. and Harada, T., 2018. Maximum classifier discrepancy for unsupervised domain adaptation.
       In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 3723-3732).
Reference code: /~https://github.com/thuml/Transfer-Learning-Library
'''
import torch
import logging
from tqdm import tqdm
import torch.nn.functional as F
from collections import defaultdict

import modules
from train_utils import TrainerBase


def classifier_discrepancy(predictions1: torch.Tensor, predictions2: torch.Tensor) -> torch.Tensor:
    return torch.mean(torch.abs(predictions1 - predictions2))


class Trainer(TrainerBase):
    
    def __init__(self, args):
        super(Trainer, self).__init__(args)
        if args.backbone == 'CNN':
            self.G = modules.MSCNN(in_channel=1).to(self.device)
        elif args.backbone == 'ResNet':
            self.G = modules.ResNet(in_channel=1, layers=[2, 2, 2, 2]).to(self.device)
        else:
            raise Exception(f"unknown backbone type {args.backbone}")
        self.C1 = modules.MLP(input_size=self.G.out_dim, output_size=args.num_classes[0],
                              dropout=args.dropout, last=None).to(self.device)
        self.C2 = modules.MLP(input_size=self.G.out_dim, output_size=args.num_classes[0],
                              dropout=args.dropout, last=None).to(self.device)
        self._init_data()

        if args.train_mode == 'single_source':
            self.src = args.source_name[0]
        elif args.train_mode == 'source_combine':
            self.src = 'concat_source'
        elif args.train_mode == 'multi_source':
            raise Exception("This model cannot be trained in multi_source mode.")
        
        self.optimizer_G = self._get_optimizer(self.G)
        self.optimizer_C = self._get_optimizer([self.C1, self.C2])
        self.lr_scheduler_G = self._get_lr_scheduler(self.optimizer_G)
        self.lr_scheduler_C = self._get_lr_scheduler(self.optimizer_C)
        self.num_iter = len(self.dataloaders[self.src])
        self.num_inner_loop = 4

    def save_model(self):
        torch.save({
            'G': self.G.state_dict(),
            'C1': self.C1.state_dict(),
            'C2': self.C2.state_dict()
            }, self.args.save_path + '.pth')
        logging.info('Model saved to {}'.format(self.args.save_path + '.pth'))
    
    def load_model(self):
        logging.info('Loading model from {}'.format(self.args.load_path))
        ckpt = torch.load(self.args.load_path)
        self.G.load_state_dict(ckpt['G'])
        self.C1.load_state_dict(ckpt['C1'])
        self.C2.load_state_dict(ckpt['C2'])

    def _set_to_train(self):
        self.G.train()
        self.C1.train()
        self.C2.train()
    
    def _set_to_eval(self):
        self.G.eval()
        self.C1.eval()
        self.C2.eval()

    def _train_one_epoch(self, epoch_acc, epoch_loss):
        for _ in tqdm(range(self.num_iter), ascii=True):
            # obtain data
            target_data, _ = self._get_next_batch('train')
            source_data, source_labels = self._get_next_batch(self.src)

            # forward Step 1
            data = torch.cat((source_data, target_data), dim=0)
            self.optimizer_G.zero_grad()
            self.optimizer_C.zero_grad()
            f = self.G(data)
            y_1 = self.C1(f)
            y_2 = self.C2(f)
            y_1, _ = y_1.chunk(2, dim=0)
            y_2, _ = y_2.chunk(2, dim=0)

            # compute loss
            loss = F.cross_entropy(y_1, source_labels) + F.cross_entropy(y_2, source_labels)
            
            # log information
            epoch_acc['Classifier 1 source train']  += self._get_accuracy(y_1, source_labels)
            epoch_acc['Classifier 2 source train']  += self._get_accuracy(y_2, source_labels)
            epoch_loss['Step 1: Source domain'] += loss

            # backward
            loss.backward()
            self.optimizer_G.step()
            self.optimizer_C.step()
            
            # forward Step 2
            self.optimizer_G.zero_grad()
            self.optimizer_C.zero_grad()
            f = self.G(data)
            y_1 = self.C1(f)
            y_2 = self.C2(f)
            y1_s, y1_t = y_1.chunk(2, dim=0)
            y2_s, y2_t = y_2.chunk(2, dim=0)
            y1_t, y2_t = F.softmax(y1_t, dim=1), F.softmax(y2_t, dim=1)

            # compute loss
            loss = F.cross_entropy(y1_s, source_labels) + F.cross_entropy(y2_s, source_labels)  \
                    - self.tradeoff[0] * classifier_discrepancy(y1_t, y2_t)
            
            # log information
            epoch_loss['Step 2: Maximize discrepancy'] += loss

            # backward
            loss.backward()
            self.optimizer_C.step()

            for _ in range(self.num_inner_loop):
                # forward Step 3
                self.optimizer_G.zero_grad()
                f = self.G(target_data)
                y_1 = self.C1(f)
                y_2 = self.C2(f)
                y1_t, y2_t = F.softmax(y_1, dim=1), F.softmax(y_2, dim=1)

                # compute loss
                loss_mcd = classifier_discrepancy(y1_t, y2_t)
                loss = self.tradeoff[1] * loss_mcd

                # log information
                epoch_loss['Step 3: Minimize discrepancy'] += loss_mcd

                # backward
                loss.backward()
                self.optimizer_G.step()
        return epoch_acc, epoch_loss
    
    def _eval(self, data, actual_labels, correct, total):
        f = self.G(data)
        y_1 = self.C1(f)
        y_2 = self.C2(f)
        pred = y_1 + y_2
        actual_pred = self._get_actual_label(pred, idx=0)
        output = self._get_accuracy(actual_pred, actual_labels, return_acc=False)
        correct['acc'] += output[0]; total['acc'] += output[1]
        if self.args.da_scenario in ['open-set', 'universal']:
            output = self._get_accuracy(actual_pred, actual_labels, return_acc=False, idx=0, mode='closed-set')
            correct['Closed-set-acc'] += output[0]; total['Closed-set-acc'] += output[1]
        return correct, total

    def train(self):
        args = self.args
        best_acc = 0.0
        best_epoch = 0
   
        for epoch in range(1, args.max_epoch+1):
            logging.info('-'*5 + 'Epoch {}/{}'.format(epoch, args.max_epoch) + '-'*5)
            
            # Update the learning rate
            if self.lr_scheduler_G is not None:
                logging.info('current lr: {}'.format(self.lr_scheduler_G.get_last_lr()))
   
            # Each epoch has a training and val phase
            epoch_acc = defaultdict(float)
   
            # Set model to train mode or evaluate mode
            self._set_to_train()
            epoch_loss = defaultdict(float)
            self.tradeoff = self._get_tradeoff(args.tradeoff, epoch) 
            
            epoch_acc, epoch_loss = self._train_one_epoch(epoch_acc, epoch_loss)
            
            # Print the train and val information via each epoch
            for key in epoch_loss.keys():
                if key == 'Step 3: Minimize discrepancy':
                    logging.info('Train-Loss {}: {:.4f}'.format(key, epoch_loss[key]/(4.*self.num_iter)))
                else:
                    logging.info('Train-Loss {}: {:.4f}'.format(key, epoch_loss[key]/self.num_iter))
            for key in epoch_acc.keys():
                logging.info('Train-Acc {}: {:.4f}'.format(key, epoch_acc[key]/self.num_iter))
            
            # Log the best model according to the val accuracy
            new_acc = self.test()
            if new_acc >= best_acc:
                best_acc = new_acc
                best_epoch = epoch
            logging.info("The best model epoch {}, val-acc {:.4f}".format(best_epoch, best_acc))
                    
            if self.lr_scheduler_G is not None:
                self.lr_scheduler_G.step()
            if self.lr_scheduler_C is not None:
                self.lr_scheduler_C.step()