train_subset.py

import argparse
import copy
import logging
import math
import os
import random
import shutil
import time
from collections import OrderedDict
from subset_selection.ssl import *
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, Subset, Dataset
from torch.utils.data.distributed import DistributedSampler
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from dataset.cifar import DATASET_GETTERS
from utils import AverageMeter, accuracy
from typing import TypeVar, Generic, Iterable, Iterator, Sequence, List, Optional, Tuple
logger = logging.getLogger(__name__)
best_acc = 0


T_co = TypeVar('T_co', covariant=True)

class CustomSubset(Dataset[T_co]):
    r"""
    Subset of a dataset at specified indices.

    Arguments:
        dataset (Dataset): The whole Dataset
        indices (sequence): Indices in the whole set selected for subset
        gammas (sequence): Weight of each data point in the subset
    """
    dataset: Dataset[T_co]
    indices: Sequence[int]

    def __init__(self, dataset: Dataset[T_co], indices: Sequence[int], gammas: Sequence[float]) -> None:
        self.dataset = dataset
        self.indices = indices
        self.gammas = gammas

    def __getitem__(self, idx):
        return self.dataset[self.indices[idx]], self.gammas[idx]

    def __len__(self):
        return len(self.indices)


def save_checkpoint(state, is_best, checkpoint, filename='checkpoint.pth.tar'):
    filepath = os.path.join(checkpoint, filename)
    torch.save(state, filepath)
    if is_best:
        shutil.copyfile(filepath, os.path.join(checkpoint,
                                               'model_best.pth.tar'))


def set_seed(args):
    random.seed(args.seed)
    np.random.seed(args.seed)
    torch.manual_seed(args.seed)
    if args.n_gpu > 0:
        torch.cuda.manual_seed_all(args.seed)


def get_cosine_schedule_with_warmup(optimizer,
                                    num_warmup_steps,
                                    num_training_steps,
                                    num_cycles=7. / 16.,
                                    last_epoch=-1):
    def _lr_lambda(current_step):
        if current_step < num_warmup_steps:
            return float(current_step) / float(max(1, num_warmup_steps))
        no_progress = float((current_step%num_training_steps) - num_warmup_steps) / \
                      float(max(1, num_training_steps - num_warmup_steps))
        return max(0., math.cos(math.pi * num_cycles * no_progress))

    return LambdaLR(optimizer, _lr_lambda, last_epoch)


def interleave(x, size):
    s = list(x.shape)
    return x.reshape([-1, size] + s[1:]).transpose(0, 1).reshape([-1] + s[1:])


def de_interleave(x, size):
    s = list(x.shape)
    return x.reshape([size, -1] + s[1:]).transpose(0, 1).reshape([-1] + s[1:])


def main():
    parser = argparse.ArgumentParser(description='PyTorch FixMatch Training')
    parser.add_argument('--gpu-id', default='0', type=int,
                        help='id(s) for CUDA_VISIBLE_DEVICES')
    parser.add_argument('--num-workers', type=int, default=4,
                        help='number of workers')
    parser.add_argument('--dataset', default='cifar10', type=str,
                        choices=['cifar10', 'cifar100'],
                        help='dataset name')
    parser.add_argument('--num-labeled', type=int, default=4000,
                        help='number of labeled data')
    parser.add_argument("--expand-labels", action="store_true",
                        help="expand labels to fit eval steps")
    parser.add_argument('--arch', default='wideresnet', type=str,
                        choices=['wideresnet', 'resnext'],
                        help='dataset name')
    parser.add_argument('--total-steps', default=2 ** 20, type=int,
                        help='number of total steps to run')
    parser.add_argument('--eval-step', default=1024, type=int,
                        help='number of eval steps to run')
    parser.add_argument('--start-epoch', default=0, type=int,
                        help='manual epoch number (useful on restarts)')
    parser.add_argument('--batch-size', default=64, type=int,
                        help='train batchsize')
    parser.add_argument('--lr', '--learning-rate', default=0.03, type=float,
                        help='initial learning rate')
    parser.add_argument('--warmup', default=0, type=float,
                        help='warmup epochs (unlabeled data based)')
    parser.add_argument('--wdecay', default=5e-4, type=float,
                        help='weight decay')
    parser.add_argument('--nesterov', action='store_true', default=True,
                        help='use nesterov momentum')
    parser.add_argument('--use-ema', action='store_true', default=True,
                        help='use EMA model')
    parser.add_argument('--ema-decay', default=0.999, type=float,
                        help='EMA decay rate')
    parser.add_argument('--mu', default=7, type=int,
                        help='coefficient of unlabeled batch size')
    parser.add_argument('--lambda-u', default=1, type=float,
                        help='coefficient of unlabeled loss')
    parser.add_argument('--T', default=1, type=float,
                        help='pseudo label temperature')
    parser.add_argument('--threshold', default=0.95, type=float,
                        help='pseudo label threshold')
    parser.add_argument('--out', default='result',
                        help='directory to output the result')
    parser.add_argument('--resume', default='', type=str,
                        help='path to latest checkpoint (default: none)')
    parser.add_argument('--seed', default=None, type=int,
                        help="random seed")
    parser.add_argument("--amp", action="store_true",
                        help="use 16-bit (mixed) precision through NVIDIA apex AMP")
    parser.add_argument("--opt_level", type=str, default="O1",
                        help="apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
                             "See details at https://nvidia.github.io/apex/amp.html")
    parser.add_argument("--local_rank", type=int, default=-1,
                        help="For distributed training: local_rank")
    parser.add_argument('--no-progress', action='store_true',
                        help="don't use progress bar")
    parser.add_argument('--dss_strategy', type=str, default='GLISTER-Warm',
                        help="Subset Selection Strategy")
    parser.add_argument('--fraction', type=float, default=0.3,
                        help="Subset Size")
    parser.add_argument('--R', type=int, default=20,
                        help="Subset selection every R epoch")
    parser.add_argument('--kappa', type=float, default=0.5,
                        help="Warm Start Coefficient")
    parser.add_argument('--max', action='store_true', default=False)
    args = parser.parse_args()
    global best_acc

    def create_model(args):
        if args.arch == 'wideresnet':
            import models.wideresnet as models
            model = models.build_wideresnet(depth=args.model_depth,
                                            widen_factor=args.model_width,
                                            dropout=0,
                                            num_classes=args.num_classes)
        elif args.arch == 'resnext':
            import models.resnext as models
            model = models.build_resnext(cardinality=args.model_cardinality,
                                         depth=args.model_depth,
                                         width=args.model_width,
                                         num_classes=args.num_classes)
        logger.info("Total params: {:.2f}M".format(
            sum(p.numel() for p in model.parameters()) / 1e6))
        return model

    if args.local_rank == -1:
        device = torch.device('cuda', args.gpu_id)
        args.world_size = 1
        args.n_gpu = torch.cuda.device_count()
    else:
        torch.cuda.set_device(args.local_rank)
        device = torch.device('cuda', args.local_rank)
        torch.distributed.init_process_group(backend='nccl')
        args.world_size = torch.distributed.get_world_size()
        args.n_gpu = 1

    args.device = device

    logging.basicConfig(
        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
        datefmt="%m/%d/%Y %H:%M:%S",
        level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)

    logger.warning(
        f"Process rank: {args.local_rank}, "
        f"device: {args.device}, "
        f"n_gpu: {args.n_gpu}, "
        f"distributed training: {bool(args.local_rank != -1)}, "
        f"16-bits training: {args.amp}", )

    logger.info(dict(args._get_kwargs()))

    if args.seed is not None:
        set_seed(args)

    if args.local_rank in [-1, 0]:
        os.makedirs(args.out, exist_ok=True)
        args.writer = SummaryWriter(args.out)

    if args.dataset == 'cifar10':
        args.num_classes = 10
        if args.arch == 'wideresnet':
            args.model_depth = 28
            args.model_width = 2
        elif args.arch == 'resnext':
            args.model_cardinality = 4
            args.model_depth = 28
            args.model_width = 4

    elif args.dataset == 'cifar100':
        args.num_classes = 100
        if args.arch == 'wideresnet':
            args.model_depth = 28
            args.model_width = 8
        elif args.arch == 'resnext':
            args.model_cardinality = 8
            args.model_depth = 29
            args.model_width = 64

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()

    labeled_dataset, unlabeled_dataset, test_dataset = DATASET_GETTERS[args.dataset](
        args, './data')

    if args.local_rank == 0:
        torch.distributed.barrier()

    if args.local_rank not in [-1, 0]:
        torch.distributed.barrier()

    model = create_model(args)
    model1 = create_model(args)
    if args.local_rank == 0:
        torch.distributed.barrier()

    model.to(args.device)
    model1.to(args.device)

    no_decay = ['bias', 'bn']
    grouped_parameters = [
        {'params': [p for n, p in model.named_parameters() if not any(
            nd in n for nd in no_decay)], 'weight_decay': args.wdecay},
        {'params': [p for n, p in model.named_parameters() if any(
            nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    optimizer = optim.SGD(grouped_parameters, lr=args.lr,
                          momentum=0.9, nesterov=args.nesterov)

    no_decay = ['bias', 'bn']
    grouped_parameters = [
        {'params': [p for n, p in model1.named_parameters() if not any(
            nd in n for nd in no_decay)], 'weight_decay': args.wdecay},
        {'params': [p for n, p in model1.named_parameters() if any(
            nd in n for nd in no_decay)], 'weight_decay': 0.0}
    ]
    optimizer1 = optim.SGD(grouped_parameters, lr=args.lr,
                          momentum=0.9, nesterov=args.nesterov)

    args.epochs = math.ceil(args.total_steps / args.eval_step)
    if args.dss_strategy == 'Full':
        scheduler = get_cosine_schedule_with_warmup(
            optimizer, args.warmup, int(args.epochs * 1024))
    else:
        scheduler = get_cosine_schedule_with_warmup(
            optimizer, args.warmup, int(args.epochs * args.fraction * 1024))

    if args.use_ema:
        from models.ema import ModelEMA
        ema_model = ModelEMA(args, model, args.ema_decay)

    args.start_epoch = 0

    if args.resume:
        logger.info("==> Resuming from checkpoint..")
        assert os.path.isfile(
            args.resume), "Error: no checkpoint directory found!"
        args.out = os.path.dirname(args.resume)
        checkpoint = torch.load(args.resume)
        best_acc = checkpoint['best_acc']
        args.start_epoch = checkpoint['epoch']
        model.load_state_dict(checkpoint['state_dict'])
        if args.use_ema:
            ema_model.ema.load_state_dict(checkpoint['ema_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])

    if args.amp:
        from apex import amp
        model, optimizer = amp.initialize(
            model, optimizer, opt_level=args.opt_level)

    if args.local_rank != -1:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[args.local_rank],
            output_device=args.local_rank, find_unused_parameters=True)

    logger.info("***** Running training *****")
    logger.info(f"  Task = {args.dataset}@{args.num_labeled}")
    logger.info(f"  Num Epochs = {args.epochs}")
    logger.info(f"  Batch size per GPU = {args.batch_size}")
    logger.info(
        f"  Total train batch size = {args.batch_size * args.world_size}")
    logger.info(f"  Total optimization steps = {args.total_steps}")

    model.zero_grad()
    train_subset(args, labeled_dataset, unlabeled_dataset, test_dataset,
          model, model1, optimizer, optimizer1, ema_model, scheduler)


def train_subset(args, labeled_dataset, unlabeled_dataset, test_dataset,
          model, model1, optimizer, optimizer1, ema_model, scheduler):
    if args.dss_strategy == 'Full':
        max_epochs = args.epochs
    else:
        if args.max:
            max_epochs = args.epochs
        else:
            max_epochs = int(args.epochs * args.fraction)

    sel_iter = int((args.R * len(unlabeled_dataset) * args.fraction) // (args.batch_size * args.mu))
    N = len(unlabeled_dataset)
    bud = int(args.fraction * N)
    start_idxs = np.random.choice(N, size=bud, replace=False)
    gammas = torch.ones(len(start_idxs))
    unlabeled_subset = CustomSubset(unlabeled_dataset, start_idxs, gammas=gammas)

    train_sampler = RandomSampler if args.local_rank == -1 else DistributedSampler

    unlabeled_subloader = DataLoader(
                            unlabeled_subset,
                            sampler=train_sampler(unlabeled_subset),
                            batch_size=args.batch_size * args.mu,
                            num_workers=args.num_workers,
                            drop_last=True)

    labeled_seq_loader = DataLoader(
        labeled_dataset,
        sampler=SequentialSampler(labeled_dataset),
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        drop_last=False)

    unlabeled_seq_loader = DataLoader(
        unlabeled_dataset,
        sampler=SequentialSampler(unlabeled_dataset),
        batch_size=args.batch_size * args.mu,
        num_workers=args.num_workers,
        drop_last=False)

    test_loader = DataLoader(
        test_dataset,
        sampler=SequentialSampler(test_dataset),
        batch_size=args.batch_size,
        num_workers=args.num_workers)

    criterion_nored = nn.CrossEntropyLoss(reduction='none')

    if args.dss_strategy == 'GradMatch':
        # OMPGradMatch Selection strategy
        setf_model = OMPGradMatchStrategy(unlabeled_seq_loader, labeled_seq_loader, model1,
                                          criterion_nored, args.lr, args.device, args.num_classes,
                                          True, 'PerClassPerGradient', args, valid=False, lam=0.25, eps=1e-10)
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'GradMatchPB':
        setf_model = OMPGradMatchStrategy(unlabeled_seq_loader, labeled_seq_loader, model1,
                                          criterion_nored, args.lr, args.device, args.num_classes,
                                          True, 'PerBatch', args, valid=False, lam=0.25, eps=1e-10)
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'GLISTER':
        # GLISTER Selection strategy
        setf_model = GLISTERStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                     args.lr, args.device, args.num_classes, False, 'Stochastic', args, r=int(bud), valid=True)
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'GLISTER_UL':
        # GLISTER Selection strategy
        setf_model = GLISTERStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                     args.lr, args.device, args.num_classes, False, 'Stochastic', args, r=int(bud), valid=False)
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'CRAIG':
        # CRAIG Selection strategy
        setf_model = CRAIGStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                    args.device, args.num_classes, False, False, 'PerClass', args, optimizer='lazy')
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'CRAIGPB':
        # CRAIG Selection strategy
        setf_model = CRAIGStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                    args.device, args.num_classes, False, False, 'PerBatch', args, optimizer='lazy')
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'CRAIG-Warm':
        # CRAIG Selection strategy
        setf_model = CRAIGStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                    args.device, args.num_classes, False, False, 'PerClass', args, optimizer='lazy')
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'CRAIGPB-Warm':
        # CRAIG Selection strategy
        setf_model = CRAIGStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                    args.device, args.num_classes, False, False, 'PerBatch', args, optimizer='lazy')
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'Random':
        # Random Selection strategy
        setf_model = RandomStrategy(unlabeled_seq_loader, online=False)
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'Random-Online':
        # Random-Online Selection strategy
        setf_model = RandomStrategy(unlabeled_seq_loader, online=True)
        kappa_epochs = int(args.kappa * args.epochs * args.fraction)
        sel_iter = int((len(unlabeled_dataset) * args.fraction) // (args.batch_size * args.mu))

    elif args.dss_strategy == 'GLISTER-Warm':
        # GLISTER Selection strategy
        setf_model = GLISTERStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                   args.lr, args.device, args.num_classes, False, 'Stochastic', args, r=int(bud), valid=True)
        kappa_epochs=int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'GLISTER_UL-Warm':
        # GLISTER Selection strategy
        setf_model = GLISTERStrategy(unlabeled_seq_loader, labeled_seq_loader, model1, criterion_nored,
                                  args.lr, args.device, args.num_classes, False, 'Stochastic', args, r=int(bud), valid=True)
        kappa_epochs=int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'GradMatch-Warm':
        # OMPGradMatch Selection strategy
        setf_model = OMPGradMatchStrategy(unlabeled_seq_loader, labeled_seq_loader, model1,
                                          criterion_nored, args.lr, args.device, args.num_classes,
                                          True, 'PerClassPerGradient', args, valid=False, lam=0.25, eps=1e-10)
        kappa_epochs=int(args.kappa * args.epochs * args.fraction)
        # full_epochs = round(kappa_epochs * self.configdata['dss_strategy']['fraction'])

    elif args.dss_strategy == 'GradMatchPB-Warm':
        # OMPGradMatch Selection strategy
        setf_model = OMPGradMatchStrategy(unlabeled_seq_loader, labeled_seq_loader, model1,
                                          criterion_nored, args.lr, args.device, args.num_classes,
                                          True, 'PerBatch', args, valid=False, lam=0.25, eps=1e-10)
        kappa_epochs=int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'Random-Warm':
        setf_model = RandomStrategy(unlabeled_seq_loader, online=False)
        kappa_epochs=int(args.kappa * args.epochs * args.fraction)

    elif args.dss_strategy == 'Full':
        kappa_epochs=int(args.kappa * args.epochs * args.fraction)

    if args.amp:
        from apex import amp
    global best_acc
    test_accs = []
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    losses_x = AverageMeter()
    losses_u = AverageMeter()
    mask_probs = AverageMeter()
    end = time.time()
    curr_iter = 0

    labeled_trainloader = DataLoader(
        labeled_dataset,
        sampler=train_sampler(labeled_dataset),
        batch_size=args.batch_size,
        num_workers=args.num_workers,
        drop_last=True)

    unlabeled_trainloader = DataLoader(
        unlabeled_dataset,
        sampler=train_sampler(unlabeled_dataset),
        batch_size=args.batch_size * args.mu,
        num_workers=args.num_workers,
        drop_last=True)

    subset_selection_time = 0
    training_time = 0

    if args.dss_strategy == 'Full':
        model.train()

        if args.world_size > 1:
            labeled_epoch = 0
            unlabeled_epoch = 0
            labeled_trainloader.sampler.set_epoch(labeled_epoch)
            unlabeled_trainloader.sampler.set_epoch(unlabeled_epoch)

        labeled_iter = iter(labeled_trainloader)
        unlabeled_iter = iter(unlabeled_trainloader)

        for epoch in range(args.start_epoch, max_epochs):
            if not args.no_progress:
                p_bar = tqdm(range(args.eval_step),
                             disable=args.local_rank not in [-1, 0])

            for batch_idx in range(args.eval_step):
                batch_start_time = time.time()
                curr_iter += 1
                try:
                    inputs_x, targets_x = labeled_iter.next()
                except:
                    if args.world_size > 1:
                        labeled_epoch += 1
                        labeled_trainloader.sampler.set_epoch(labeled_epoch)
                    labeled_iter = iter(labeled_trainloader)
                    inputs_x, targets_x = labeled_iter.next()

                try:
                    (inputs_u_w, inputs_u_s), _ = unlabeled_iter.next()
                except:
                    if args.world_size > 1:
                        unlabeled_epoch += 1
                        unlabeled_trainloader.sampler.set_epoch(unlabeled_epoch)
                    unlabeled_iter = iter(unlabeled_trainloader)
                    (inputs_u_w, inputs_u_s), _ = unlabeled_iter.next()

                data_time.update(time.time() - end)
                batch_size = inputs_x.shape[0]
                inputs = interleave(
                    torch.cat((inputs_x, inputs_u_w, inputs_u_s)), 2 * args.mu + 1).to(args.device)
                targets_x = targets_x.to(args.device)
                logits = model(inputs)
                logits = de_interleave(logits, 2 * args.mu + 1)
                logits_x = logits[:batch_size]
                logits_u_w, logits_u_s = logits[batch_size:].chunk(2)
                del logits

                Lx = F.cross_entropy(logits_x, targets_x, reduction='mean')

                pseudo_label = torch.softmax(logits_u_w.detach() / args.T, dim=-1)
                max_probs, targets_u = torch.max(pseudo_label, dim=-1)
                mask = max_probs.ge(args.threshold).float()

                Lu = (F.cross_entropy(logits_u_s, targets_u,
                                      reduction='none') * mask).mean()

                loss = Lx + args.lambda_u * Lu

                if args.amp:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()

                losses.update(loss.item())
                losses_x.update(Lx.item())
                losses_u.update(Lu.item())
                optimizer.step()
                scheduler.step()
                if args.use_ema:
                    ema_model.update(model)
                training_time += (time.time() - batch_start_time)
                model.zero_grad()
                batch_time.update(time.time() - end)
                end = time.time()
                mask_probs.update(mask.mean().item())
                if not args.no_progress:
                    p_bar.set_description(
                        "Train Epoch: {epoch}/{epochs:4}. Iter: {batch:4}/{iter:4}. LR: {lr:.4f}. Data: {data:.3f}s. Batch: {bt:.3f}s. Loss: {loss:.4f}. Loss_x: {loss_x:.4f}. Loss_u: {loss_u:.4f}. Mask: {mask:.2f}. ".format(
                            epoch=epoch + 1,
                            epochs=args.epochs,
                            batch=batch_idx + 1,
                            iter=args.eval_step,
                            lr=scheduler.get_last_lr()[0],
                            data=data_time.avg,
                            bt=batch_time.avg,
                            loss=losses.avg,
                            loss_x=losses_x.avg,
                            loss_u=losses_u.avg,
                            mask=mask_probs.avg))
                    p_bar.update()

            if not args.no_progress:
                p_bar.close()

            if args.use_ema:
                test_model = ema_model.ema
            else:
                test_model = model

            if args.local_rank in [-1, 0]:
                test_loss, test_acc = test(args, test_loader, test_model, epoch)

                args.writer.add_scalar('train/1.train_loss', losses.avg, epoch)
                args.writer.add_scalar('train/2.train_loss_x', losses_x.avg, epoch)
                args.writer.add_scalar('train/3.train_loss_u', losses_u.avg, epoch)
                args.writer.add_scalar('train/4.mask', mask_probs.avg, epoch)
                args.writer.add_scalar('test/1.test_acc', test_acc, epoch)
                args.writer.add_scalar('test/2.test_loss', test_loss, epoch)

                is_best = test_acc > best_acc
                best_acc = max(test_acc, best_acc)

                model_to_save = model.module if hasattr(model, "module") else model
                if args.use_ema:
                    ema_to_save = ema_model.ema.module if hasattr(
                        ema_model.ema, "module") else ema_model.ema
                save_checkpoint({
                    'epoch': epoch + 1,
                    'state_dict': model_to_save.state_dict(),
                    'ema_state_dict': ema_to_save.state_dict() if args.use_ema else None,
                    'acc': test_acc,
                    'best_acc': best_acc,
                    'optimizer': optimizer.state_dict(),
                    'scheduler': scheduler.state_dict(),
                }, is_best, args.out)

                test_accs.append(test_acc)
                logger.info('Best top-1 acc: {:.2f}'.format(best_acc))
                logger.info('Mean top-1 acc: {:.2f}\n'.format(
                    np.mean(test_accs[-20:])))

    elif args.dss_strategy in ['CRAIG', 'CRAIGPB', 'GradMatch', 'GradMatchPB', 'GLISTER', 'GLISTER_UL', 'Random', 'Random-Online']:
        model.train()
        if args.world_size > 1:
            labeled_epoch = 0
            unlabeled_epoch = 0
            labeled_trainloader.sampler.set_epoch(labeled_epoch)
            unlabeled_subloader.sampler.set_epoch(unlabeled_epoch)

        labeled_iter = iter(labeled_trainloader)
        unlabeled_iter = iter(unlabeled_subloader)

        for epoch in range(args.start_epoch, max_epochs):
            if not args.no_progress:
                p_bar = tqdm(range(args.eval_step),
                             disable=args.local_rank not in [-1, 0])

            for batch_idx in range(args.eval_step):
                batch_start_time = time.time()
                curr_iter += 1
                try:
                    inputs_x, targets_x = labeled_iter.next()
                except:
                    if args.world_size > 1:
                        labeled_epoch += 1
                        labeled_trainloader.sampler.set_epoch(labeled_epoch)
                    labeled_iter = iter(labeled_trainloader)
                    inputs_x, targets_x = labeled_iter.next()

                try:
                    ((inputs_u_w, inputs_u_s), _), weights = unlabeled_iter.next()
                except:
                    if args.world_size > 1:
                        unlabeled_epoch += 1
                        unlabeled_subloader.sampler.set_epoch(unlabeled_epoch)
                    unlabeled_iter = iter(unlabeled_subloader)
                    ((inputs_u_w, inputs_u_s), _), weights = unlabeled_iter.next()

                data_time.update(time.time() - end)
                weights = weights.to(args.device)
                batch_size = inputs_x.shape[0]
                inputs = interleave(
                    torch.cat((inputs_x, inputs_u_w, inputs_u_s)), 2 * args.mu + 1).to(args.device)
                targets_x = targets_x.to(args.device)
                logits = model(inputs)
                logits = de_interleave(logits, 2 * args.mu + 1)
                logits_x = logits[:batch_size]
                logits_u_w, logits_u_s = logits[batch_size:].chunk(2)
                del logits

                Lx = F.cross_entropy(logits_x, targets_x, reduction='mean')

                pseudo_label = torch.softmax(logits_u_w.detach() / args.T, dim=-1)
                max_probs, targets_u = torch.max(pseudo_label, dim=-1)
                mask = max_probs.ge(args.threshold).float()

                Lu = (F.cross_entropy(logits_u_s, targets_u,
                                      reduction='none') * mask * weights).mean()
                loss = Lx + args.lambda_u * Lu

                if args.amp:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()

                losses.update(loss.item())
                losses_x.update(Lx.item())
                losses_u.update(Lu.item())
                optimizer.step()
                scheduler.step()
                if args.use_ema:
                    ema_model.update(model)
                training_time += (time.time() - batch_start_time)
                model.zero_grad()
                batch_time.update(time.time() - end)
                end = time.time()
                mask_probs.update(mask.mean().item())
                if not args.no_progress:
                    p_bar.set_description(
                        "Train Epoch: {epoch}/{epochs:4}. Iter: {batch:4}/{iter:4}. LR: {lr:.4f}. Data: {data:.3f}s. Batch: {bt:.3f}s. Loss: {loss:.4f}. Loss_x: {loss_x:.4f}. Loss_u: {loss_u:.4f}. Mask: {mask:.2f}. ".format(
                            epoch=epoch + 1,
                            epochs=args.epochs,
                            batch=batch_idx + 1,
                            iter=args.eval_step,
                            lr=scheduler.get_last_lr()[0],
                            data=data_time.avg,
                            bt=batch_time.avg,
                            loss=losses.avg,
                            loss_x=losses_x.avg,
                            loss_u=losses_u.avg,
                            mask=mask_probs.avg))
                    p_bar.update()

                if curr_iter % sel_iter == 0:
                    start_time = time.time()
                    # Perform Subset Selection
                    cached_state_dict = copy.deepcopy(model.state_dict())
                    clone_dict = copy.deepcopy(model.state_dict())
                    subset_idxs, gammas = setf_model.select(int(bud), clone_dict)
                    model.load_state_dict(cached_state_dict)
                    idxs = subset_idxs
                    if args.dss_strategy in ['GradMatch', 'GradMatchPB', 'CRAIG', 'CRAIGPB']:
                        gammas = torch.from_numpy(np.array(gammas)).to(torch.float32)
                    subset_selection_time += (time.time() - start_time)
                    unlabeled_subset = CustomSubset(unlabeled_dataset, idxs, gammas)
                    unlabeled_subloader = DataLoader(
                        unlabeled_subset,
                        sampler=train_sampler(unlabeled_subset),
                        batch_size=args.batch_size * args.mu,
                        num_workers=args.num_workers,
                        drop_last=True)
                    unlabeled_iter = iter(unlabeled_subloader)
            if not args.no_progress:
                p_bar.close()

            if args.use_ema:
                test_model = ema_model.ema
            else:
                test_model = model

            if args.local_rank in [-1, 0]:
                test_loss, test_acc = test(args, test_loader, test_model, epoch)

                args.writer.add_scalar('train/1.train_loss', losses.avg, epoch)
                args.writer.add_scalar('train/2.train_loss_x', losses_x.avg, epoch)
                args.writer.add_scalar('train/3.train_loss_u', losses_u.avg, epoch)
                args.writer.add_scalar('train/4.mask', mask_probs.avg, epoch)
                args.writer.add_scalar('test/1.test_acc', test_acc, epoch)
                args.writer.add_scalar('test/2.test_loss', test_loss, epoch)

                is_best = test_acc > best_acc
                best_acc = max(test_acc, best_acc)

                model_to_save = model.module if hasattr(model, "module") else model
                if args.use_ema:
                    ema_to_save = ema_model.ema.module if hasattr(
                        ema_model.ema, "module") else ema_model.ema
                save_checkpoint({
                    'epoch': epoch + 1,
                    'state_dict': model_to_save.state_dict(),
                    'ema_state_dict': ema_to_save.state_dict() if args.use_ema else None,
                    'acc': test_acc,
                    'best_acc': best_acc,
                    'optimizer': optimizer.state_dict(),
                    'scheduler': scheduler.state_dict(),
                }, is_best, args.out)

                test_accs.append(test_acc)
                logger.info('Best top-1 acc: {:.2f}'.format(best_acc))
                logger.info('Mean top-1 acc: {:.2f}\n'.format(
                    np.mean(test_accs[-20:])))

    elif args.dss_strategy in ['CRAIGPB-Warm', 'CRAIG-Warm', 'GradMatch-Warm', 'GradMatchPB-Warm',
                               'GLISTER-Warm', 'GLISTER_UL-Warm', 'Random-Warm']:
        model.train()

        if args.world_size > 1:
            labeled_epoch = 0
            unlabeled_epoch = 0
            labeled_trainloader.sampler.set_epoch(labeled_epoch)
            unlabeled_trainloader.sampler.set_epoch(unlabeled_epoch)

        labeled_iter = iter(labeled_trainloader)
        unlabeled_iter = iter(unlabeled_trainloader)

        for epoch in range(args.start_epoch, kappa_epochs):
            #if args.world_size > 1:
            if not args.no_progress:
                p_bar = tqdm(range(args.eval_step),
                             disable=args.local_rank not in [-1, 0])

            for batch_idx in range(args.eval_step):
                batch_start_time = time.time()
                curr_iter += 1
                try:
                    inputs_x, targets_x = labeled_iter.next()
                except:
                    if args.world_size > 1:
                        labeled_epoch += 1
                        labeled_trainloader.sampler.set_epoch(labeled_epoch)
                    labeled_iter = iter(labeled_trainloader)
                    inputs_x, targets_x = labeled_iter.next()

                try:
                    (inputs_u_w, inputs_u_s), _ = unlabeled_iter.next()
                except:
                    if args.world_size > 1:
                        unlabeled_epoch += 1
                        unlabeled_subloader.sampler.set_epoch(unlabeled_epoch)
                    unlabeled_iter = iter(unlabeled_trainloader)
                    (inputs_u_w, inputs_u_s), _ = unlabeled_iter.next()

                data_time.update(time.time() - end)
                batch_size = inputs_x.shape[0]

                inputs = interleave(
                    torch.cat((inputs_x, inputs_u_w, inputs_u_s)), 2 * args.mu + 1).to(args.device)

                targets_x = targets_x.to(args.device)
                logits = model(inputs)
                logits = de_interleave(logits, 2 * args.mu + 1)
                logits_x = logits[:batch_size]
                logits_u_w, logits_u_s = logits[batch_size:].chunk(2)
                del logits
                Lx = F.cross_entropy(logits_x, targets_x, reduction='mean')
                pseudo_label = torch.softmax(logits_u_w.detach() / args.T, dim=-1)
                max_probs, targets_u = torch.max(pseudo_label, dim=-1)
                mask = max_probs.ge(args.threshold).float()
                Lu = (F.cross_entropy(logits_u_s, targets_u,
                                      reduction='none') * mask).mean()
                loss = Lx + args.lambda_u * Lu
                if args.amp:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()

                losses.update(loss.item())
                losses_x.update(Lx.item())
                losses_u.update(Lu.item())
                optimizer.step()
                scheduler.step()
                if args.use_ema:
                    ema_model.update(model)
                training_time += (time.time() - batch_start_time)
                model.zero_grad()
                batch_time.update(time.time() - end)
                end = time.time()
                mask_probs.update(mask.mean().item())
                if not args.no_progress:
                    p_bar.set_description(
                        "Train Epoch: {epoch}/{epochs:4}. Iter: {batch:4}/{iter:4}. LR: {lr:.4f}. Data: {data:.3f}s. Batch: {bt:.3f}s. Loss: {loss:.4f}. Loss_x: {loss_x:.4f}. Loss_u: {loss_u:.4f}. Mask: {mask:.2f}. ".format(
                            epoch=epoch + 1,
                            epochs=args.epochs,
                            batch=batch_idx + 1,
                            iter=args.eval_step,
                            lr=scheduler.get_last_lr()[0],
                            data=data_time.avg,
                            bt=batch_time.avg,
                            loss=losses.avg,
                            loss_x=losses_x.avg,
                            loss_u=losses_u.avg,
                            mask=mask_probs.avg))
                    p_bar.update()
            if not args.no_progress:
                p_bar.close()

            if args.use_ema:
                test_model = ema_model.ema
            else:
                test_model = model

            if args.local_rank in [-1, 0]:
                test_loss, test_acc = test(args, test_loader, test_model, epoch)

                args.writer.add_scalar('train/1.train_loss', losses.avg, epoch)
                args.writer.add_scalar('train/2.train_loss_x', losses_x.avg, epoch)
                args.writer.add_scalar('train/3.train_loss_u', losses_u.avg, epoch)
                args.writer.add_scalar('train/4.mask', mask_probs.avg, epoch)
                args.writer.add_scalar('test/1.test_acc', test_acc, epoch)
                args.writer.add_scalar('test/2.test_loss', test_loss, epoch)

                is_best = test_acc > best_acc
                best_acc = max(test_acc, best_acc)

                model_to_save = model.module if hasattr(model, "module") else model
                if args.use_ema:
                    ema_to_save = ema_model.ema.module if hasattr(
                        ema_model.ema, "module") else ema_model.ema
                save_checkpoint({
                    'epoch': epoch + 1,
                    'state_dict': model_to_save.state_dict(),
                    'ema_state_dict': ema_to_save.state_dict() if args.use_ema else None,
                    'acc': test_acc,
                    'best_acc': best_acc,
                    'optimizer': optimizer.state_dict(),
                    'scheduler': scheduler.state_dict(),
                }, is_best, args.out)

                test_accs.append(test_acc)
                logger.info('Best top-1 acc: {:.2f}'.format(best_acc))
                logger.info('Mean top-1 acc: {:.2f}\n'.format(
                    np.mean(test_accs[-20:])))

        model.train()
        if args.world_size > 1:
            labeled_epoch = 0
            unlabeled_epoch = 0
            labeled_trainloader.sampler.set_epoch(labeled_epoch)
            unlabeled_subloader.sampler.set_epoch(unlabeled_epoch)

        labeled_iter = iter(labeled_trainloader)
        unlabeled_iter = iter(unlabeled_subloader)

        for epoch in range(kappa_epochs, max_epochs):
            #if args.world_size > 1:
            if not args.no_progress:
                p_bar = tqdm(range(args.eval_step),
                             disable=args.local_rank not in [-1, 0])

            for batch_idx in range(args.eval_step):
                batch_start_time = time.time()
                curr_iter += 1
                try:
                    inputs_x, targets_x = labeled_iter.next()
                except:
                    if args.world_size > 1:
                        labeled_epoch += 1
                        labeled_trainloader.sampler.set_epoch(labeled_epoch)
                    labeled_iter = iter(labeled_trainloader)
                    inputs_x, targets_x = labeled_iter.next()

                try:
                    ((inputs_u_w, inputs_u_s), _), weights = unlabeled_iter.next()
                except:
                    if args.world_size > 1:
                        unlabeled_epoch += 1
                        unlabeled_subloader.sampler.set_epoch(unlabeled_epoch)
                    unlabeled_iter = iter(unlabeled_subloader)
                    ((inputs_u_w, inputs_u_s), _), weights = unlabeled_iter.next()

                data_time.update(time.time() - end)
                weights = weights.to(args.device)
                batch_size = inputs_x.shape[0]
                inputs = interleave(
                    torch.cat((inputs_x, inputs_u_w, inputs_u_s)), 2 * args.mu + 1).to(args.device)
                targets_x = targets_x.to(args.device)
                logits = model(inputs)
                logits = de_interleave(logits, 2 * args.mu + 1)
                logits_x = logits[:batch_size]
                logits_u_w, logits_u_s = logits[batch_size:].chunk(2)
                del logits

                Lx = F.cross_entropy(logits_x, targets_x, reduction='mean')

                pseudo_label = torch.softmax(logits_u_w.detach() / args.T, dim=-1)
                max_probs, targets_u = torch.max(pseudo_label, dim=-1)
                mask = max_probs.ge(args.threshold).float()

                Lu = (F.cross_entropy(logits_u_s, targets_u,
                                      reduction='none') * mask * weights).mean()
                loss = Lx + args.lambda_u * Lu

                if args.amp:
                    with amp.scale_loss(loss, optimizer) as scaled_loss:
                        scaled_loss.backward()
                else:
                    loss.backward()

                losses.update(loss.item())
                losses_x.update(Lx.item())
                losses_u.update(Lu.item())
                optimizer.step()
                scheduler.step()
                if args.use_ema:
                    ema_model.update(model)
                training_time += (time.time() - batch_start_time)
                model.zero_grad()
                batch_time.update(time.time() - end)
                end = time.time()
                mask_probs.update(mask.mean().item())
                if not args.no_progress:
                    p_bar.set_description(
                        "Train Epoch: {epoch}/{epochs:4}. Iter: {batch:4}/{iter:4}. LR: {lr:.4f}. Data: {data:.3f}s. Batch: {bt:.3f}s. Loss: {loss:.4f}. Loss_x: {loss_x:.4f}. Loss_u: {loss_u:.4f}. Mask: {mask:.2f}. ".format(
                            epoch=epoch + 1,
                            epochs=args.epochs,
                            batch=batch_idx + 1,
                            iter=args.eval_step,
                            lr=scheduler.get_last_lr()[0],
                            data=data_time.avg,
                            bt=batch_time.avg,
                            loss=losses.avg,
                            loss_x=losses_x.avg,
                            loss_u=losses_u.avg,
                            mask=mask_probs.avg))
                    p_bar.update()

                if curr_iter % sel_iter == 0:
                    start_time = time.time()
                    # Perform Subset Selection
                    cached_state_dict = copy.deepcopy(model.state_dict())
                    clone_dict = copy.deepcopy(model.state_dict())
                    subset_idxs, gammas = setf_model.select(int(bud), clone_dict)
                    model.load_state_dict(cached_state_dict)
                    idxs = subset_idxs
                    if args.dss_strategy in ['GradMatch-Warm', 'GradMatchPB-Warm', 'CRAIG-Warm', 'CRAIGPB-Warm']:
                        gammas = torch.from_numpy(np.array(gammas)).to(torch.float32)
                    subset_selection_time += (time.time() - start_time)
                    unlabeled_subset = CustomSubset(unlabeled_dataset, idxs, gammas)
                    unlabeled_subloader = DataLoader(
                        unlabeled_subset,
                        sampler=train_sampler(unlabeled_subset),
                        batch_size=args.batch_size * args.mu,
                        num_workers=args.num_workers,
                        drop_last=True)
                    unlabeled_iter = iter(unlabeled_subloader)
            if not args.no_progress:
                p_bar.close()

            if args.use_ema:
                test_model = ema_model.ema
            else:
                test_model = model

            if args.local_rank in [-1, 0]:
                test_loss, test_acc = test(args, test_loader, test_model, epoch)

                args.writer.add_scalar('train/1.train_loss', losses.avg, epoch)
                args.writer.add_scalar('train/2.train_loss_x', losses_x.avg, epoch)
                args.writer.add_scalar('train/3.train_loss_u', losses_u.avg, epoch)
                args.writer.add_scalar('train/4.mask', mask_probs.avg, epoch)
                args.writer.add_scalar('test/1.test_acc', test_acc, epoch)
                args.writer.add_scalar('test/2.test_loss', test_loss, epoch)

                is_best = test_acc > best_acc
                best_acc = max(test_acc, best_acc)

                model_to_save = model.module if hasattr(model, "module") else model
                if args.use_ema:
                    ema_to_save = ema_model.ema.module if hasattr(
                        ema_model.ema, "module") else ema_model.ema
                save_checkpoint({
                    'epoch': epoch + 1,
                    'state_dict': model_to_save.state_dict(),
                    'ema_state_dict': ema_to_save.state_dict() if args.use_ema else None,
                    'acc': test_acc,
                    'best_acc': best_acc,
                    'optimizer': optimizer.state_dict(),
                    'scheduler': scheduler.state_dict(),
                }, is_best, args.out)
                test_accs.append(test_acc)
                logger.info('Best top-1 acc: {:.2f}'.format(best_acc))
                logger.info('Mean top-1 acc: {:.2f}\n'.format(np.mean(test_accs[-20:])))
                logger.info('Training Time\n'.format(subset_selection_time + training_time))

    if args.local_rank in [-1, 0]:
        args.writer.close()


def test(args, test_loader, model, epoch):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()
    end = time.time()

    if not args.no_progress:
        test_loader = tqdm(test_loader,
                           disable=args.local_rank not in [-1, 0])

    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(test_loader):
            data_time.update(time.time() - end)
            model.eval()

            inputs = inputs.to(args.device)
            targets = targets.to(args.device)
            outputs = model(inputs)
            loss = F.cross_entropy(outputs, targets)

            prec1, prec5 = accuracy(outputs, targets, topk=(1, 5))
            losses.update(loss.item(), inputs.shape[0])
            top1.update(prec1.item(), inputs.shape[0])
            top5.update(prec5.item(), inputs.shape[0])
            batch_time.update(time.time() - end)
            end = time.time()
            if not args.no_progress:
                test_loader.set_description(
                    "Test Iter: {batch:4}/{iter:4}. Data: {data:.3f}s. Batch: {bt:.3f}s. Loss: {loss:.4f}. top1: {top1:.2f}. top5: {top5:.2f}. ".format(
                        batch=batch_idx + 1,
                        iter=len(test_loader),
                        data=data_time.avg,
                        bt=batch_time.avg,
                        loss=losses.avg,
                        top1=top1.avg,
                        top5=top5.avg,
                    ))
        if not args.no_progress:
            test_loader.close()

    logger.info("top-1 acc: {:.2f}".format(top1.avg))
    logger.info("top-5 acc: {:.2f}".format(top5.avg))
    return losses.avg, top1.avg


if __name__ == '__main__':
    main()