main.py

import parser

from fastai.losses import LabelSmoothingCrossEntropy
from fastai.vision import *
import re
import argparse
import os
import shutil
import time
import math
import logging

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.backends.cudnn as cudnn
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler
import torchvision.datasets

from mean_teacher import architectures, datasets, data, losses, ramps, cli
from mean_teacher.run_context import RunContext
from mean_teacher.data import NO_LABEL
from mean_teacher.utils import *
from mean_teacher.losses import SmoothLabelCritierion
LOG = logging.getLogger('main')

args = None
best_prec1 = 0
global_step = 0
best_pred =0

def reduce_loss(loss,
                reduction='mean'): return loss.mean() if reduction == 'mean' else loss.sum() if reduction == 'sum' else loss


def linear_combination(x, y, epsilon): return epsilon * x + (1 - epsilon) * y


class LSRO(nn.Module):
    def __init__(self, epsilon: float = 0.1, reduction='mean'):
        super().__init__()
        self.epsilon = epsilon
        self.reduction = reduction

    def forward(self, preds, target):
        n = preds.size()[-1]
        log_preds = F.log_softmax(preds, dim=-1)
        loss = reduce_loss(-log_preds.sum(dim=-1), self.reduction)
        nll = F.nll_loss(F.log_softmax(log_preds, 1), target, ignore_index=-1, reduction=self.reduction)
        return linear_combination(loss / n, nll, self.epsilon)


def main(context):
    global global_step
    global best_prec1
    global best_pred
    checkpoint_path = context.transient_dir
    training_log = context.create_train_log("training")
    validation_log = context.create_train_log("validation")
    ema_validation_log = context.create_train_log("ema_validation")

    dataset_config = datasets.__dict__[args.dataset]()
    num_classes = dataset_config.pop('num_classes')
    train_loader, eval_loader, train_loader_gan = create_data_loaders(**dataset_config, args=args)

    def create_model(ema=False):
        LOG.info("=> creating {pretrained}{ema}model '{arch}'".format(
            pretrained='pre-trained ' if args.pretrained else '',
            ema='EMA ' if ema else '',
            arch=args.arch))

        model_factory = architectures.__dict__[args.arch]
        model_params = dict(pretrained=args.pretrained, num_classes=num_classes)
        model = model_factory(**model_params)
        model = nn.DataParallel(model).cuda()

        if ema:
            for param in model.parameters():
                param.detach_()

        return model

    model = create_model()
    ema_model = create_model(ema=True)

    LOG.info(parameters_string(model))

    optimizer = torch.optim.SGD(model.parameters(), args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay,
                                nesterov=args.nesterov)

    # optionally resume from a checkpoint
    if args.resume:
        assert os.path.isfile(args.resume), "=> no checkpoint found at '{}'".format(args.resume)
        LOG.info("=> loading checkpoint '{}'".format(args.resume))
        checkpoint = torch.load(args.resume)
        args.start_epoch = checkpoint['epoch']
        global_step = checkpoint['global_step']
        best_prec1 = checkpoint['best_prec1']
        model.load_state_dict(checkpoint['state_dict'])
        ema_model.load_state_dict(checkpoint['ema_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        LOG.info("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))

    cudnn.benchmark = True

    if args.evaluate:
        LOG.info("Evaluating the primary model:")
        validate(eval_loader, model, validation_log, global_step, args.start_epoch)
        LOG.info("Evaluating the EMA model:")
        validate(eval_loader, ema_model, ema_validation_log, global_step, args.start_epoch)
        return

    for epoch in range(args.start_epoch, args.epochs):
        start_time = time.time()
        # train for one epoch
        train(train_loader, train_loader_gan, model, ema_model, optimizer, epoch, training_log)
        # train(train_loader, eval_loader, model, ema_model, optimizer, epoch, training_log)
        LOG.info("--- training epoch in %s seconds ---" % (time.time() - start_time))

        if args.evaluation_epochs and (epoch + 1) % args.evaluation_epochs == 0:
            start_time = time.time()
            LOG.info("Evaluating the primary model:")
            prec1,_,_ = validate(eval_loader, model, validation_log, global_step, epoch + 1)
            LOG.info("Evaluating the EMA model:")
            ema_prec1, pred, target = validate(eval_loader, ema_model, ema_validation_log, global_step, epoch + 1)
            LOG.info("--- validation in %s seconds ---" % (time.time() - start_time))
            is_best = ema_prec1 > best_prec1
            if is_best:
                best_prec1 = max(ema_prec1, best_prec1)
                best_pred = pred
        else:
            is_best = False

        if args.checkpoint_epochs and (epoch + 1) % args.checkpoint_epochs == 0:
            save_checkpoint({
                'epoch': epoch + 1,
                'global_step': global_step,
                'arch': args.arch,
                'state_dict': model.state_dict(),
                'ema_state_dict': ema_model.state_dict(),
                'best_prec1': best_prec1,
                'best_pred': best_pred,
                'target': target,
                'optimizer': optimizer.state_dict(),
            }, is_best, checkpoint_path, epoch + 1)


def parse_dict_args(**kwargs):
    global args

    def to_cmdline_kwarg(key, value):
        if len(key) == 1:
            key = "-{}".format(key)
        else:
            key = "--{}".format(re.sub(r"_", "-", key))
        value = str(value)
        return key, value

    kwargs_pairs = (to_cmdline_kwarg(key, value)
                    for key, value in kwargs.items())
    cmdline_args = list(sum(kwargs_pairs, ()))
    args = parser.parse_args(cmdline_args)


def create_data_loaders(train_transformation,
                        eval_transformation,
                        datadir,
                        args):
    traindir = os.path.join(datadir, args.train_subdir)
    # traindir = 'E:/Pang/mean-teacher/pytorch/data-local/images/ourData/ourData200x30/by-image/leave_user2\\train'
    evaldir = os.path.join(datadir, args.eval_subdir)
    # traindir_gan = 'E:/Pang/mean-teacher/pytorch/data-local/images/ourData/ourData200x30/by-image/leave_user2/train_gan'
    traindir_gan = os.path.join(datadir, 'train_gan')
    assert_exactly_one([args.exclude_unlabeled, args.labeled_batch_size])

    dataset = torchvision.datasets.ImageFolder(traindir, train_transformation)  # 实现数据导入
    dataset_gan = torchvision.datasets.ImageFolder(traindir_gan, train_transformation)

    if args.labels:
        with open(args.labels) as f:
            labels = dict(line.split(' ') for line in f.read().splitlines())
        labeled_idxs, unlabeled_idxs = data.relabel_dataset(dataset, labels)

    if args.labels:
        with open(args.labels) as f:
            labels = dict(line.split(' ') for line in f.read().splitlines())
        labeled_idxs_gan, unlabeled_idxs_gan = data.relabel_dataset(dataset_gan, labels)


    batch_sampler = data.TwoStreamBatchSampler(
        unlabeled_idxs, labeled_idxs, args.batch_size, args.labeled_batch_size)

    args.labeled_batch_size_gan = 99
    batch_sampler_gan = data.TwoStreamBatchSampler(
        unlabeled_idxs_gan, labeled_idxs_gan, args.batch_size, args.labeled_batch_size_gan)

    train_loader = torch.utils.data.DataLoader(dataset,
                                               batch_sampler=batch_sampler,
                                               num_workers=0,
                                               pin_memory=True)

    train_loader_gan = torch.utils.data.DataLoader(dataset_gan,
                                                   batch_sampler=batch_sampler_gan,
                                                   num_workers=0,
                                                   pin_memory=True)

    eval_loader = torch.utils.data.DataLoader(
        torchvision.datasets.ImageFolder(evaldir, eval_transformation),
        batch_size=250,
        shuffle=False,
        num_workers=0,
        pin_memory=True,
        drop_last=False)

    return train_loader, eval_loader, train_loader_gan


def update_ema_variables(model, ema_model, alpha, global_step):
    # Use the true average until the exponential average is more correct
    alpha = min(1 - 1 / (global_step + 1), alpha)
    for ema_param, param in zip(ema_model.parameters(), model.parameters()):
        ema_param.data.mul_(alpha).add_(1 - alpha, param.data)


def train(train_loader, train_loader_gan, model, ema_model, optimizer, epoch, log):
    global global_step

    class_criterion = SmoothLabelCritierion(label_smoothing=0).cuda()
    class_criterion_unlabel = SmoothLabelCritierion(label_smoothing=1).cuda()


    if args.consistency_type == 'mse':
        consistency_criterion = losses.softmax_mse_loss
    elif args.consistency_type == 'kl':
        consistency_criterion = losses.softmax_kl_loss
    else:
        assert False, args.consistency_type
    residual_logit_criterion = losses.symmetric_mse_loss

    meters = AverageMeterSet()

    # switch to train mode
    model.train()
    ema_model.train()

    end = time.time()
    for i, ((input, ema_input), target) in enumerate(train_loader):
        dataiter = iter(train_loader_gan)
        (input_gan, ema_input_gan), target_gan = dataiter.next()
        # for j, ((input_gan, ema_input_gan), target_gan) in enumerate(train_loader_gan):
        # measure data loading time
        meters.update('data_time', time.time() - end)

        adjust_learning_rate(optimizer, epoch, i, len(train_loader))
        meters.update('lr', optimizer.param_groups[0]['lr'])
        # adjust_learning_rate(optimizer, epoch, i, len(train_loader_gan))
        # meters.update('lr', optimizer.param_groups[0]['lr'])

        input_var = torch.autograd.Variable(input,
                                            volatile=True)
        ema_input_var = torch.autograd.Variable(ema_input, volatile=True)
        # ema_input_var = torch.autograd.Variable(ema_input)
        target_var = torch.autograd.Variable(target.cuda())

        input_var_gan = torch.autograd.Variable(input_gan,
                                                volatile=True)
        ema_input_var_gan = torch.autograd.Variable(ema_input_gan, volatile=True)
        # ema_input_var = torch.autograd.Variable(ema_input)
        target_var_gan = torch.autograd.Variable(target_gan.cuda())

        minibatch_size = len(target_var)
        minibatch_size_gan = len(target_var_gan)
        labeled_minibatch_size = target_var.data.ne(NO_LABEL).sum()
        labeled_minibatch_size_gan = target_var_gan.data.ne(NO_LABEL).sum()
        assert labeled_minibatch_size > 0
        assert labeled_minibatch_size_gan > 0

        unlabel_minibatch_size =minibatch_size-labeled_minibatch_size
        unlabel_minibatch_size_gan =minibatch_size_gan-labeled_minibatch_size_gan
        meters.update('labeled_minibatch_size', labeled_minibatch_size)
        meters.update('labeled_minibatch_size_gan', labeled_minibatch_size_gan)

        ema_model_out = ema_model(ema_input_var)
        ema_model_out_gan = ema_model(ema_input_var_gan)
        model_out = model(input_var)
        model_out_gan = model(input_var_gan)

        if isinstance(model_out, Variable):
            assert args.logit_distance_cost < 0
            logit1 = model_out
            ema_logit = ema_model_out
        else:
            assert len(model_out) == 2
            assert len(ema_model_out) == 2
            logit1, logit2 = model_out
            ema_logit, _ = ema_model_out


        logit1_gan, logit2_gan = model_out_gan
        ema_logit_gan, _ = ema_model_out_gan

        ema_logit = Variable(ema_logit.detach().data, requires_grad=False)
        ema_logit_gan = Variable(ema_logit_gan.detach().data, requires_grad=False)

        if args.logit_distance_cost >= 0:  # -1
            class_logit, cons_logit = logit1, logit2
            res_loss = args.logit_distance_cost * residual_logit_criterion(class_logit, cons_logit) / minibatch_size
            meters.update('res_loss', res_loss.item())
        else:
            class_logit, cons_logit = logit1, logit1
            res_loss = 0

        if args.logit_distance_cost >= 0:
            class_logit_gan, cons_logit_gan = logit1_gan, logit2_gan
            res_loss_gan = args.logit_distance_cost * residual_logit_criterion(class_logit, cons_logit) / minibatch_size
            meters.update('res_loss', res_loss.item())
        else:
            class_logit_gan, cons_logit_gan = logit1_gan, logit1_gan
            res_loss_gan = 0

        class_loss = class_criterion(class_logit, target_var) / labeled_minibatch_size
        unlabel_class_logit = class_logit[:unlabel_minibatch_size]
        unlabel_target_var = target_var[:unlabel_minibatch_size]

        unlabel_class_loss = class_criterion_unlabel(unlabel_class_logit, unlabel_target_var) / unlabel_minibatch_size
        unlabel_class_logit_gan = class_logit_gan[:minibatch_size - labeled_minibatch_size]
        unlabel_target_var_gan = target_var_gan[:minibatch_size - labeled_minibatch_size]
        unlabel_class_loss_gan = class_criterion_unlabel(unlabel_class_logit_gan, unlabel_target_var_gan) / unlabel_minibatch_size_gan

        # meters.update('class_loss', class_loss.data[0])
        meters.update('class_loss', class_loss.item())

        ema_class_loss = class_criterion(ema_logit, target_var) / minibatch_size

        ema_class_loss_gan = class_criterion_unlabel(ema_logit_gan, target_var_gan) / minibatch_size_gan
        # meters.update('ema_class_loss', ema_class_loss.data[0])
        meters.update('ema_class_loss', ema_class_loss.item())
        meters.update('ema_class_loss_gan', ema_class_loss_gan.item())

        if args.consistency:
            consistency_weight = get_current_consistency_weight(epoch)
            meters.update('cons_weight', consistency_weight)
            consistency_loss = consistency_weight * consistency_criterion(cons_logit, ema_logit) / minibatch_size
            meters.update('cons_loss', consistency_loss.item())
        else:
            consistency_loss = 0
            meters.update('cons_loss', 0)

        if args.consistency:
            consistency_weight_gan = get_current_consistency_weight(epoch)
            meters.update('cons_weight_gan', consistency_weight_gan)
            consistency_loss_gan = consistency_weight_gan * consistency_criterion(cons_logit_gan,
                                                                                  ema_logit_gan) / minibatch_size
            meters.update('cons_loss_gan', consistency_loss_gan.item())
        else:
            consistency_loss = 0
            meters.update('cons_loss', 0)
        class_loss_all = 0.98*class_loss + 0.01*unlabel_class_loss + 0.01*unlabel_class_loss_gan
        consistency_loss_all = 0.99 * consistency_loss + 0.01 * consistency_loss_gan
        # class_loss_all = class_loss
        # consistency_loss_all = consistency_loss
        loss = class_loss_all + consistency_loss_all + res_loss
        assert not (np.isnan(loss.item()) or loss.item() > 1e5), 'Loss explosion: {}'.format(loss.item())
        meters.update('loss', loss.item())

        prec1, prec5 = accuracy(class_logit.data, target_var.data, topk=(1, 5))
        meters.update('top1', prec1[0], labeled_minibatch_size)
        meters.update('error1', 100. - prec1[0], labeled_minibatch_size)
        meters.update('top5', prec5[0], labeled_minibatch_size)
        meters.update('error5', 100. - prec5[0], labeled_minibatch_size)

        ema_prec1, ema_prec5 = accuracy(ema_logit.data, target_var.data, topk=(1, 5))
        meters.update('ema_top1', ema_prec1[0], labeled_minibatch_size)
        meters.update('ema_error1', 100. - ema_prec1[0], labeled_minibatch_size)
        meters.update('ema_top5', ema_prec5[0], labeled_minibatch_size)
        meters.update('ema_error5', 100. - ema_prec5[0], labeled_minibatch_size)

        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        global_step += 1
        update_ema_variables(model, ema_model, args.ema_decay, global_step)

        # measure elapsed time
        meters.update('batch_time', time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            LOG.info(
                'Epoch: [{0}][{1}/{2}]\t'
                'Time {meters[batch_time]:.3f}\t'
                'Data {meters[data_time]:.3f}\t'
                'Class {meters[class_loss]:.4f}\t'
                'Cons {meters[cons_loss]:.4f}\t'
                'Prec@1 {meters[top1]:.3f}\t'
                'Prec@5 {meters[top5]:.3f}'.format(
                    epoch, i, len(train_loader), meters=meters))
            log.record(epoch + i / len(train_loader), {
                'step': global_step,
                **meters.values(),
                **meters.averages(),
                **meters.sums()
            })


def validate(eval_loader, model, log, global_step, epoch):
    class_criterion = nn.CrossEntropyLoss(size_average=False, ignore_index=NO_LABEL).cuda()
    meters = AverageMeterSet()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (input, target) in enumerate(eval_loader):
        meters.update('data_time', time.time() - end)

        # input_var = torch.autograd.Variable(input, volatile=True)
        # target_var = torch.autograd.Variable(target.cuda(async=True), volatile=True)
        with torch.no_grad():
            input_var = torch.autograd.Variable(input)
        with torch.no_grad():
            target_var = torch.autograd.Variable(target.cuda())

        minibatch_size = len(target_var)
        labeled_minibatch_size = target_var.data.ne(NO_LABEL).sum()
        assert labeled_minibatch_size > 0
        meters.update('labeled_minibatch_size', labeled_minibatch_size)

        # compute output
        output1, output2 = model(input_var)
        softmax1, softmax2 = F.softmax(output1, dim=1), F.softmax(output2, dim=1)
        class_loss = class_criterion(output1, target_var) / minibatch_size

        # measure accuracy and record loss
        _, pred = output1.topk(1, 1, True, True)
        prec1, prec5 = accuracy(output1.data, target_var.data, topk=(1, 5))
        meters.update('class_loss', class_loss.item(), labeled_minibatch_size)
        meters.update('top1', prec1[0], labeled_minibatch_size)
        meters.update('error1', 100.0 - prec1[0], labeled_minibatch_size)
        meters.update('top5', prec5[0], labeled_minibatch_size)
        meters.update('error5', 100.0 - prec5[0], labeled_minibatch_size)

        # measure elapsed time
        meters.update('batch_time', time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            LOG.info(
                'Test: [{0}/{1}]\t'
                'Time {meters[batch_time]:.3f}\t'
                'Data {meters[data_time]:.3f}\t'
                'Class {meters[class_loss]:.4f}\t'
                'Prec@1 {meters[top1]:.3f}\t'
                'Prec@5 {meters[top5]:.3f}'.format(
                    i, len(eval_loader), meters=meters))

    LOG.info(' * Prec@1 {top1.avg:.3f}\tPrec@5 {top5.avg:.3f}'
             .format(top1=meters['top1'], top5=meters['top5']))
    log.record(epoch, {
        'step': global_step,
        **meters.values(),
        **meters.averages(),
        **meters.sums()
    })

    return meters['top1'].avg, pred, target_var.data


def save_checkpoint(state, is_best, dirpath, epoch):
    filename = 'checkpoint.{}.ckpt'.format(epoch)
    checkpoint_path = os.path.join(dirpath, filename)
    best_path = os.path.join(dirpath, 'best.ckpt')
    torch.save(state, checkpoint_path)
    LOG.info("--- checkpoint saved to %s ---" % checkpoint_path)
    if is_best:
        shutil.copyfile(checkpoint_path, best_path)
        LOG.info("--- checkpoint copied to %s ---" % best_path)


def adjust_learning_rate(optimizer, epoch, step_in_epoch, total_steps_in_epoch):
    lr = args.lr
    epoch = epoch + step_in_epoch / total_steps_in_epoch

    # LR warm-up to handle large minibatch sizes from https://arxiv.org/abs/1706.02677
    lr = ramps.linear_rampup(epoch, args.lr_rampup) * (args.lr - args.initial_lr) + args.initial_lr

    # Cosine LR rampdown from https://arxiv.org/abs/1608.03983 (but one cycle only)
    if args.lr_rampdown_epochs:
        assert args.lr_rampdown_epochs >= args.epochs
        lr *= ramps.cosine_rampdown(epoch, args.lr_rampdown_epochs)

    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


def get_current_consistency_weight(epoch):
    # Consistency ramp-up from https://arxiv.org/abs/1610.02242
    return args.consistency * ramps.sigmoid_rampup(epoch, args.consistency_rampup)


def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    labeled_minibatch_size = max(target.ne(NO_LABEL).sum(), 1e-8)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()  # 转置
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:  # k=1,k=5
        correct_k = correct[:k].contiguous().view(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / labeled_minibatch_size))
    return res


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    args = cli.parse_commandline_args()
    main(RunContext(__file__, 0))