Add FKD support

FKD (https://arxiv.org/abs/2112.01528)

FKD_train.py

@@ -0,0 +1,296 @@
+#!/usr/bin/env python3
+"""Script to train a model through soft labels on ImageNet's train set."""
+
+import argparse
+import logging
+import pprint
+import os
+import sys
+import time
+import math
+import numpy as np
+
+import torch
+from torch import nn
+
+from loss import discriminatorLoss
+
+import imagenet
+from models import model_factory
+from models import discriminator
+import opts
+import test
+import utils
+from utils_FKD import Recover_soft_label
+
+def parse_args(argv):
+    """Parse arguments @argv and return the flags needed for training."""
+    parser = argparse.ArgumentParser(description=__doc__, allow_abbrev=False)
+
+    group = parser.add_argument_group('General Options')
+    opts.add_general_flags(group)
+
+    group = parser.add_argument_group('Dataset Options')
+    opts.add_dataset_flags(group)
+
+    group = parser.add_argument_group('Model Options')
+    opts.add_model_flags(group)
+
+    group = parser.add_argument_group('Soft Label Options')
+    opts.add_teacher_flags(group)
+
+    group = parser.add_argument_group('Training Options')
+    opts.add_training_flags(group)
+
+    group = parser.add_argument_group('CutMix Training Options')
+    opts.add_cutmix_training_flags(group)
+
+    args = parser.parse_args(argv)
+
+    return args
+
+
+class LearningRateRegime:
+    """Encapsulates the learning rate regime for training a model.
+
+    Args:
+        @intervals (list): A list of triples (start, end, lr). The intervals
+            are inclusive (for start <= epoch <= end, lr will be used). The
+            start of each interval must be right after the end of its previous
+            interval.
+    """
+
+    def __init__(self, regime):
+        if len(regime) % 3 != 0:
+            raise ValueError("Regime length should be devisible by 3.")
+        intervals = list(zip(regime[0::3], regime[1::3], regime[2::3]))
+        self._validate_intervals(intervals)
+        self.intervals = intervals
+        self.num_epochs = intervals[-1][1]
+
+    @classmethod
+    def _validate_intervals(cls, intervals):
+        if type(intervals) is not list:
+            raise TypeError("Intervals must be a list of triples.")
+        elif len(intervals) == 0:
+            raise ValueError("Intervals must be a non empty list.")
+        elif intervals[0][0] != 1:
+            raise ValueError("Intervals must start from 1: {}".format(intervals))
+        elif any(end < start for (start, end, lr) in intervals):
+            raise ValueError("End of intervals must be greater or equal than their"
+                             " start: {}".format(intervals))
+        elif any(intervals[i][1] + 1 != intervals[i + 1][0]
+                 for i in range(len(intervals) - 1)):
+            raise ValueError("Start of each each interval must be the end of its "
+                             "previous interval plus one: {}".format(intervals))
+
+    def get_lr(self, epoch):
+        for (start, end, lr) in self.intervals:
+            if start <= epoch <= end:
+                return lr
+        raise ValueError("Invalid epoch {} for regime {!r}".format(
+            epoch, self.intervals))
+
+
+def _set_learning_rate(optimizer, lr):
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Decay the learning rate based on schedule"""
+    lr = args.lr
+    if args.cos:  # cosine lr schedule
+        lr *= 0.5 * (1. + math.cos(math.pi * epoch / args.epochs))
+    else:  # stepwise lr schedule
+        for milestone in args.schedule:
+            milestone = math.ceil(milestone / args.num_crops)
+            lr *= 0.1 if epoch >= milestone else 1.
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+def _get_learning_rate(optimizer):
+    return max(param_group['lr'] for param_group in optimizer.param_groups)
+
+
+def train_for_one_epoch(model, g_loss, discriminator_loss, train_loader, optimizer, epoch_number, args):
+    model.train()
+    g_loss.train()
+
+    data_time_meter = utils.AverageMeter()
+    batch_time_meter = utils.AverageMeter()
+    g_loss_meter = utils.AverageMeter(recent=100)
+    d_loss_meter = utils.AverageMeter(recent=100)
+    top1_meter = utils.AverageMeter(recent=100)
+    top5_meter = utils.AverageMeter(recent=100)
+
+    timestamp = time.time()
+    for i, (images, labels, soft_labels) in enumerate(train_loader):
+        batch_size = args.batch_size
+
+        # Record data time
+        data_time_meter.update(time.time() - timestamp)
+
+        images = torch.cat(images, dim=0)
+        soft_labels = torch.cat(soft_labels, dim=0)
+        labels = torch.cat(labels, dim=0)
+
+        if args.soft_label_type == 'ori':
+            soft_labels = soft_labels.cuda()
+        else:
+            soft_labels = Recover_soft_label(soft_labels, args.soft_label_type, args.num_classes)
+            soft_labels = soft_labels.cuda()
+
+        if utils.is_model_cuda(model):
+            images = images.cuda()
+            labels = labels.cuda()
+
+        if args.w_cutmix == True:
+            r = np.random.rand(1)
+            if args.beta > 0 and r < args.cutmix_prob:
+                # generate mixed sample
+                lam = np.random.beta(args.beta, args.beta)
+                rand_index = torch.randperm(images.size()[0]).cuda()
+                target_a = soft_labels
+                target_b = soft_labels[rand_index]
+                bbx1, bby1, bbx2, bby2 = utils.rand_bbox(images.size(), lam)
+                images[:, :, bbx1:bbx2, bby1:bby2] = images[rand_index, :, bbx1:bbx2, bby1:bby2]
+                lam = 1 - ((bbx2 - bbx1) * (bby2 - bby1) / (images.size()[-1] * images.size()[-2]))
+
+        # Forward pass, backward pass, and update parameters.
+        output = model(images)
+        # output, soft_label, soft_no_softmax = outputs
+        if args.w_cutmix == True:
+            g_loss_output1 = g_loss((output, target_a), labels)
+            g_loss_output2 = g_loss((output, target_b), labels)
+        else:
+            g_loss_output = g_loss((output, soft_labels), labels)
+        if args.use_discriminator_loss:
+            # Our stored label is "after softmax", this is slightly different from original MEAL V2 
+            # that used probibilaties "before softmax" for the discriminator.
+            if args.w_cutmix == True:
+                d_loss_value = discriminator_loss([output], [target_a]) * lam + discriminator_loss([output], [target_b]) * (1 - lam)
+            else:
+                d_loss_value = discriminator_loss([output], [soft_labels])
+
+        # Sometimes loss function returns a modified version of the output,
+        # which must be used to compute the model accuracy.
+        if args.w_cutmix == True:
+            if isinstance(g_loss_output1, tuple):
+                g_loss_value1, output1 = g_loss_output1
+                g_loss_value2, output2 = g_loss_output2
+                g_loss_value = g_loss_value1 * lam + g_loss_value2 * (1 - lam)
+            else:
+                g_loss_value = g_loss_output1 * lam + g_loss_output2 * (1 - lam)
+        else:
+            if isinstance(g_loss_output, tuple):
+                g_loss_value, output = g_loss_output
+            else:
+                g_loss_value = g_loss_output
+
+        if args.use_discriminator_loss:
+            loss_value = g_loss_value + d_loss_value
+        else:
+            loss_value = g_loss_value 
+
+        loss_value.backward()
+
+        # Update parameters and reset gradients.
+        optimizer.step()
+        optimizer.zero_grad()
+
+        # Record loss and model accuracy.
+        g_loss_meter.update(g_loss_value.item(), batch_size)
+        d_loss_meter.update(d_loss_value.item(), batch_size)
+
+        top1, top5 = utils.topk_accuracy(output, labels, recalls=(1, 5))
+        top1_meter.update(top1, batch_size)
+        top5_meter.update(top5, batch_size)
+
+        # Record batch time
+        batch_time_meter.update(time.time() - timestamp)
+        timestamp = time.time()
+
+        if i%20 == 0:
+            logging.info(
+                'Epoch: [{epoch}][{batch}/{epoch_size}]\t'
+                'Time {batch_time.value:.2f} ({batch_time.average:.2f})   '
+                'Data {data_time.value:.2f} ({data_time.average:.2f})   '
+                'G_Loss {g_loss.value:.3f} {{{g_loss.average:.3f}, {g_loss.average_recent:.3f}}}    '
+                'D_Loss {d_loss.value:.3f} {{{d_loss.average:.3f}, {d_loss.average_recent:.3f}}}    '
+                'Top-1 {top1.value:.2f} {{{top1.average:.2f}, {top1.average_recent:.2f}}}    '
+                'Top-5 {top5.value:.2f} {{{top5.average:.2f}, {top5.average_recent:.2f}}}    '
+                'LR {lr:.5f}'.format(
+                    epoch=epoch_number, batch=i + 1, epoch_size=len(train_loader),
+                    batch_time=batch_time_meter, data_time=data_time_meter,
+                    g_loss=g_loss_meter, d_loss=d_loss_meter, top1=top1_meter, top5=top5_meter,
+                    lr=_get_learning_rate(optimizer)))
+    # Log the overall train stats
+    logging.info(
+        'Epoch: [{epoch}] -- TRAINING SUMMARY\t'
+        'Time {batch_time.sum:.2f}   '
+        'Data {data_time.sum:.2f}   '
+        'G_Loss {g_loss.average:.3f}     '
+        'D_Loss {d_loss.average:.3f}     '
+        'Top-1 {top1.average:.2f}    '
+        'Top-5 {top5.average:.2f}    '.format(
+            epoch=epoch_number, batch_time=batch_time_meter, data_time=data_time_meter,
+            g_loss=g_loss_meter, d_loss=d_loss_meter, top1=top1_meter, top5=top5_meter))
+
+
+def save_checkpoint(checkpoints_dir, model, optimizer, epoch):
+    model_state_file = os.path.join(checkpoints_dir, 'model_state_{:02}.pytar'.format(epoch))
+    optim_state_file = os.path.join(checkpoints_dir, 'optim_state_{:02}.pytar'.format(epoch))
+    torch.save(model.state_dict(), model_state_file)
+    torch.save(optimizer.state_dict(), optim_state_file)
+
+
+def create_optimizer(model,  discriminator_parameters, momentum=0.9, weight_decay=0):
+    # Get model parameters that require a gradient.
+    parameters = [{'params': model.parameters()}, discriminator_parameters]
+    optimizer = torch.optim.SGD(parameters, lr=0,
+                                momentum=momentum, weight_decay=weight_decay)
+    return optimizer
+
+def create_discriminator_criterion(args):
+    d = discriminator.Discriminator(outputs_size=1000, K=8).cuda()
+    d = torch.nn.DataParallel(d)
+    update_parameters = {'params': d.parameters(), "lr": args.d_lr}
+    discriminators_criterion = discriminatorLoss(d).cuda()
+    if len(args.gpus) > 1:
+        discriminators_criterion = torch.nn.DataParallel(discriminators_criterion, device_ids=args.gpus)
+    return discriminators_criterion, update_parameters
+
+def main(argv):
+    """Run the training script with command line arguments @argv."""
+    args = parse_args(argv)
+    utils.general_setup(args.save, args.gpus)
+
+    logging.info("Arguments parsed.\n{}".format(pprint.pformat(vars(args))))
+
+    # convert to TRUE number of loading-images since we use multiple crops from the same image within a minbatch
+    args.batch_size = math.ceil(args.batch_size / args.num_crops)
+
+    # Create the train and the validation data loaders.
+    train_loader = imagenet.get_train_loader_FKD(args.imagenet, args.batch_size,
+                                             args.num_workers, args.image_size, args.num_crops, args.softlabel_path)
+    val_loader = imagenet.get_val_loader(args.imagenet, args.batch_size,
+                                         args.num_workers, args.image_size)
+    # Create model with optional teachers.
+    model, loss = model_factory.create_model(
+        args.model, args.student_state_file, args.gpus, args.teacher_model,
+        args.teacher_state_file, True)
+    logging.info("Model:\n{}".format(model))
+
+    discriminator_loss, update_parameters = create_discriminator_criterion(args)
+
+    optimizer = create_optimizer(model, update_parameters, args.momentum, args.weight_decay)
+
+    for epoch in range(args.start_epoch, args.epochs, args.num_crops):
+        adjust_learning_rate(optimizer, epoch, args)
+        train_for_one_epoch(model, loss, discriminator_loss, train_loader, optimizer, epoch, args)
+        test.test_for_one_epoch(model, loss, val_loader, epoch)
+        save_checkpoint(args.save, model, optimizer, epoch)
+
+
+if __name__ == '__main__':
+    main(sys.argv[1:])

README.md

@@ -23,6 +23,24 @@ If you find our code is helpful for your research, please cite:
 
 ## News
 
+**[Dec. 5, 2021]** **New:** Add [FKD](https://arxiv.org/abs/2112.01528) training support. We highly recommend to use FKD for training MEAL V2 models, which will be 2~4x faster while reaching similar accuracy. 
+
+- Download our [soft label](http://zhiqiangshen.com/projects/FKD/index.html) for MEAL V2.
+- run script/FKD_train.sh with the desired model architecture and the path to the ImageNet dataset and soft label, for example:
+
+  ```shell
+  # 224 x 224 ResNet-50
+  python FKD_train.sh --save MEAL_V2_resnet50_224 \
+  --batch-size 512 -j 48 \
+  --model resnet50 --epochs 180 \
+  --teacher-model gluon_senet154,gluon_resnet152_v1s \
+  --imagenet [imagenet-folder with train and val folders] \
+  --num_crops 8 --soft_label_type marginal_smoothing_k5 \
+  --softlabel_path [path of soft label] \
+  --schedule 100 240 --use-discriminator-loss
+  ```
+Add `--cos` if you would like to train with cosine learning rate.
+
 **New:** Basically, adding back tricks (cosine *lr*, etc.) into MEAL V2 can consistently improve the accuracy:
 
 <div align=center>

imagenet.py

@@ -6,13 +6,33 @@
 from torchvision import datasets as torch_datasets
 from torchvision import transforms
 
+from utils_FKD import RandomResizedCrop_FKD,RandomHorizontalFlip_FKD,ImageFolder_FKD,Compose_FKD
+from torchvision.transforms import InterpolationMode
 
 def get_train_loader(imagenet_path, batch_size, num_workers, image_size):
     train_dataset = ImageNet(imagenet_path, image_size, is_train=True)
     return data_utils.DataLoader(
         train_dataset, shuffle=True, batch_size=batch_size, pin_memory=True,
         num_workers=num_workers)
 
+def get_train_loader_FKD(imagenet_path, batch_size, num_workers, image_size, num_crops, softlabel_path):
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+    train_dataset = ImageFolder_FKD(
+        num_crops=num_crops,
+        softlabel_path=softlabel_path,
+        root=os.path.join(imagenet_path, 'train'),
+        transform=Compose_FKD(transforms=[
+            RandomResizedCrop_FKD(size=224,
+                                  scale=(0.08, 1),
+                                  interpolation='bilinear'), 
+            RandomHorizontalFlip_FKD(),
+            transforms.ToTensor(),
+            normalize,
+        ]))
+    return data_utils.DataLoader(
+        train_dataset, shuffle=True, batch_size=batch_size, pin_memory=True,
+        num_workers=num_workers)
 
 def get_val_loader(imagenet_path, batch_size, num_workers, image_size):
     val_dataset = ImageNet(imagenet_path, image_size, is_train=False)
@@ -50,3 +70,5 @@ def __init__(self, root_dir, im_size, is_train):
                 transforms.Normalize(ImageNet.MEAN, ImageNet.STD),
             ])
         super().__init__(root_dir, transform=transform)
+
+

models/model_factory.py

@@ -59,17 +59,21 @@ def teachers(teachers=['resnet50'], state_file=None):
 
 
 def create_model(model_name, student_state_file=None, gpus=[], teacher=None,
-                 teacher_state_file=None):
-    # model = _create_model(model_name)
-    model = _create_checkpoint_model(model_name, student_state_file)
-    model.LR_REGIME = [1, 100, 0.01, 101, 300, 0.001] # LR_REGIME 
-    if teacher is not None:
-        # assert teacher_state_file is not None, "Teacher state is None."
-
-        teacher = teachers(teacher.split(","), teacher_state_file)
-        model = teacher_wrapper.ModelDistillationWrapper(model, teacher)
+                 teacher_state_file=None, FKD=True):
+    if FKD:
+        model = _create_checkpoint_model(model_name, student_state_file)
         loss = kd_loss.KLLoss()
+        return model, loss
     else:
-        loss = nn.CrossEntropyLoss()
-
-    return model, loss
+        model = _create_checkpoint_model(model_name, student_state_file)
+        model.LR_REGIME = [0, 100, 0.01, 101, 300, 0.001] # LR_REGIME 
+        if teacher is not None:
+            # assert teacher_state_file is not None, "Teacher state is None."
+
+            teacher = teachers(teacher.split(","), teacher_state_file)
+            model = teacher_wrapper.ModelDistillationWrapper(model, teacher)
+            loss = kd_loss.KLLoss()
+        else:
+            loss = nn.CrossEntropyLoss()
+
+        return model, loss