train_ofa_net.py

"""
Trains the OFAKWSNet.
Please refer to the README.md or the report for more information.
"""

import argparse
import os
import random
import time

import numpy as np
import torch

from once_for_all.elastic_nn.networks.ofa_kws_net import OFAKWSNet
from once_for_all.elastic_nn.training.progressive_shrinking import load_models
from once_for_all.networks.kws_net import KWSNetLarge
from once_for_all.run_manager.run_config import KWSRunConfig
from once_for_all.run_manager.run_manager import RunManager
from utils.config_utils import set_ft_extr_params_to_args

parser = argparse.ArgumentParser()
parser.add_argument(
    "--task",
    type=str,
    default="normal",
    choices=[
        "normal",
        "kernel",
        "depth",
        "expand",
        "width"
    ],
)
parser.add_argument("--phase", type=int, default=1, choices=[1, 2, 3])
parser.add_argument("--resume", action="store_true")
parser.add_argument("--ft_extr_type",
                    type=str,
                    default="mfcc",
                    choices=[
                        "mfcc",
                        "mel_spectrogram",
                        "log_mel_spectrogram",
                        "spectrogram",
                        "log_spectrogram",
                        "linear_stft",
                        "raw"
                    ]
                    )

parser.add_argument("--params_id", type=int,  default=1)

args = parser.parse_args()

args.path = "exp/" + args.ft_extr_type + str(args.params_id)

args.kd_ratio = .3
args.width_mult_list = "1.0"

if args.task == "normal":
    args.path += "/normal"
    args.dynamic_batch_size = 1
    args.n_epochs = 140
    args.base_lr = 1e-3
    args.warmup_epochs = 5
    args.warmup_lr = -1
    args.ks_list = "7"
    args.depth_list = "4"
    args.expand_list = "3"
    args.kd_ratio = 0
elif args.task == "kernel":
    args.path += "/normal2kernel"
    args.dynamic_batch_size = 1
    args.n_epochs = 100
    args.base_lr = 1e-3
    args.warmup_epochs = 5
    args.warmup_lr = -1
    args.ks_list = "3,5,7"
    args.depth_list = "4"
    args.expand_list = "3"

elif args.task == "depth":
    args.path += "/kernel2kernel_depth/phase%d" % args.phase
    args.dynamic_batch_size = 2
    if args.phase == 1:
        args.n_epochs = 25
        args.base_lr = 1e-3
        args.warmup_epochs = 0
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "3,4"
        args.expand_list = "3"

    elif args.phase == 2:
        args.n_epochs = 25
        args.base_lr = 1e-3
        args.warmup_epochs = 5
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "2,3,4"
        args.expand_list = "3"

    else:
        args.n_epochs = 100
        args.base_lr = 1e-3
        args.warmup_epochs = 5
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "1,2,3,4"
        args.expand_list = "3"


elif args.task == "expand":
    args.path += "/kernel_depth2kernel_depth_expand/phase%d" % args.phase
    args.dynamic_batch_size = 4
    if args.phase == 1:
        args.n_epochs = 25
        args.base_lr = 1e-3
        args.warmup_epochs = 0
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "1,2,3,4"
        args.expand_list = "2,3"

    elif args.phase == 2:
        args.n_epochs = 100
        args.base_lr = 1e-3
        args.warmup_epochs = 5
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "1,2,3,4"
        args.expand_list = "1,2,3"

elif args.task == "width":  # Unused in our implementation of OFAKWS
    args.path += "/kernel_depth2kernel_depth_expand_width/phase%d" % args.phase
    args.dynamic_batch_size = 8
    if args.phase == 1:
        args.n_epochs = 25  # 25
        args.base_lr = 1e-3
        args.warmup_epochs = 0
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "1,2,3,4"
        args.expand_list = "1,2,3"
        args.width_mult_list = "0.75,1.0"

    elif args.phase == 2:
        args.n_epochs = 100  # 55 # 120
        args.base_lr = 1e-3
        args.warmup_epochs = 5
        args.warmup_lr = -1
        args.ks_list = "3,5,7"
        args.depth_list = "1,2,3,4"
        args.expand_list = "1,2,3"
        args.width_mult_list = "0.5,0.75,1.0"
else:
    raise NotImplementedError

args.manual_seed = 0
args.base_batch_size = 256  # 512
args.valid_size = 10000
args.momentum = 0.9
args.no_nesterov = False
args.weight_decay = 3e-5
args.label_smoothing = 0.1
args.no_decay_keys = "bn#bias"
args.fp16_allreduce = False
args.model_init = "he_fout"
args.validation_frequency = 3
args.print_frequency = 5
args.n_worker = 8
args.bn_momentum = 0.1
args.bn_eps = 1e-5
args.dropout = 0.1
args.kd_type = "ce"

"""Set ft_extr_params_list depending on the ft_extr_type"""

args = set_ft_extr_params_to_args(args)

if __name__ == "__main__":
    os.makedirs(args.path, exist_ok=True)
    start = time.time()

    num_gpus = torch.cuda.device_count()
    print("Using %f gpus" % num_gpus)

    # Set random seed
    torch.manual_seed(args.manual_seed)
    torch.cuda.manual_seed_all(args.manual_seed)
    np.random.seed(args.manual_seed)
    random.seed(args.manual_seed)

    # Cuda Setup
    if torch.cuda.is_available():
        # Pin GPU to be used to process local rank (one GPU per process)
        torch.backends.cudnn.enabled = True
        torch.backends.cudnn.benchmark = True
        torch.cuda.manual_seed(args.manual_seed)
        print('Using GPU.')
    else:
        print('Using CPU.')

    args.teacher_path = "exp/" + args.ft_extr_type + str(args.params_id)
    args.teacher_path += "/normal/checkpoint/checkpoint.pth.tar"

    # build run config from args
    args.lr_schedule_param = None
    args.lr_schedule_type = "cosine"
    args.opt_type = "adam"  # cosine, sgd, adam

    args.opt_param = {
        "momentum": args.momentum,
        "nesterov": not args.no_nesterov,
    }
    args.init_lr = args.base_lr * num_gpus  # linearly rescale the learning rate
    if args.warmup_lr < 0:
        args.warmup_lr = args.base_lr

    args.train_batch_size = args.base_batch_size
    args.test_batch_size = args.base_batch_size * 4

    run_config = KWSRunConfig(**args.__dict__, num_replicas=num_gpus)

    # Print run config information
    print("Run config:")
    for k, v in run_config.config.items():
        print("\t%s: %s" % (k, v))

    # Build net from args
    args.ks_list = [int(ks) for ks in args.ks_list.split(",")]
    args.depth_list = [int(d) for d in args.depth_list.split(",")]
    args.expand_list = [int(e) for e in args.expand_list.split(",")]
    args.width_mult_list = [float(w) for w in args.width_mult_list.split(",")]

    # Instantiate OFA KWS model
    net = OFAKWSNet(
        n_classes=12,
        bn_param=(args.bn_momentum, args.bn_eps),
        dropout_rate=args.dropout,
        width_mult_list=args.width_mult_list,
        ks_list=args.ks_list,
        depth_list=args.depth_list,
        expand_ratio_list=args.expand_list
    )

    """ RunManager """
    run_manager = RunManager(
        args.path,
        net,
        run_config
    )
    run_manager.save_config()

    # Instantiate largest KWS model possible
    if args.kd_ratio > 0:
        args.teacher_model = KWSNetLarge(
            n_classes=12,
            bn_param=(args.bn_momentum, args.bn_eps),
            dropout_rate=0,
            width_mult=max(args.width_mult_list),
            ks=max(args.ks_list),
            depth=max(args.depth_list),
            expand_ratio=max(args.expand_list)
        )

        if torch.cuda.is_available():
            args.teacher_model.cuda()

        #  load teacher net weights
        load_models(run_manager, args.teacher_model, model_path=args.teacher_path)

        # Validate teacher net
        teach_validate_func_dict = {
            "ft_extr_type": args.ft_extr_type,
            "ft_extr_params_list": args.ft_extr_params_list,
            "width_mult_list": [max(args.width_mult_list)],
            "ks_list": [max(args.ks_list)],
            "depth_list": [max(net.depth_list)],
            "expand_list": [max(args.expand_list)],
        }
        print("Teacher validation feature extraction type: ", teach_validate_func_dict['ft_extr_type'])
        print("Teacher validation feature extraction parameter search space: ",
              teach_validate_func_dict['ft_extr_params_list'])
        run_manager.validate_all_resolution(is_test=True, net=args.teacher_model)

    """Training"""
    from once_for_all.elastic_nn.training.progressive_shrinking import (
        validate,
        train,
    )

    validate_func_dict = {
        "ft_extr_type": args.ft_extr_type,
        "ft_extr_params_list": {args.ft_extr_params_list[0], args.ft_extr_params_list[-1]} if isinstance(args.ft_extr_params_list, list) else args.ft_extr_params_list,
        "width_mult_list": sorted({min(args.width_mult_list), max(args.width_mult_list)}),
        "ks_list": sorted({min(args.ks_list), max(args.ks_list)}),
        "depth_list": sorted({max(min(net.depth_list), 1), max(net.depth_list)}),
        # no reason to test net with all depths = 0
        "expand_list": sorted({min(args.expand_list), max(args.expand_list)}),
    }
    print("Validation feature extraction type: ", validate_func_dict['ft_extr_type'])
    print("Validation feature extraction parameter search space: ", validate_func_dict['ft_extr_params_list'])

    args.ofa_checkpoint_path = "exp/" + args.ft_extr_type + str(args.params_id)
    if args.task == "normal":
        if args.resume:
            args.ofa_checkpoint_path += "/normal/checkpoint/model_best.pth.tar"
            load_models(
                run_manager,
                run_manager.net,
                args.ofa_checkpoint_path,
            )
            run_manager.write_log(
                "%.3f\t%.3f\t%.3f\t%s"
                % validate(run_manager, is_test=True, **validate_func_dict),
                "valid",
            )
            print("Resuming training ")
        print("Start large net training")
        train(
            run_manager,
            args,
            lambda _run_manager, epoch, is_test: validate(
                _run_manager, epoch, is_test, **validate_func_dict
            ))

    elif args.task == "kernel":
        validate_func_dict["ks_list"] = sorted(args.ks_list)
        if run_manager.start_epoch == 0:
            args.ofa_checkpoint_path += "/normal/checkpoint/model_best.pth.tar"

            load_models(
                run_manager,
                run_manager.net,
                args.ofa_checkpoint_path,
            )
            run_manager.write_log(
                "%.3f\t%.3f\t%.3f\t%s"
                % validate(run_manager, is_test=True, **validate_func_dict),
                "valid",
            )
            print("Start elastic kernel training")
        else:
            assert args.resume
        train(
            run_manager,
            args,
            lambda _run_manager, epoch, is_test: validate(
                _run_manager, epoch, is_test, **validate_func_dict
            ),
        )

    elif args.task == "depth":
        from once_for_all.elastic_nn.training.progressive_shrinking import (
            train_elastic_depth,
        )

        if args.phase == 1:
            args.ofa_checkpoint_path += "/normal2kernel/checkpoint/model_best.pth.tar"
        elif args.phase == 2:
            args.ofa_checkpoint_path += "/kernel2kernel_depth/phase1/checkpoint/model_best.pth.tar"
        else:
            args.ofa_checkpoint_path += "/kernel2kernel_depth/phase2/checkpoint/model_best.pth.tar"

        print("Start elastic depth training")

        train_elastic_depth(train, run_manager, args, validate_func_dict)

    elif args.task == "expand":
        from once_for_all.elastic_nn.training.progressive_shrinking import (
            train_elastic_expand,
        )

        if args.phase == 1:
            args.ofa_checkpoint_path += "/kernel2kernel_depth/phase2/checkpoint/model_best.pth.tar"
        elif args.phase == 2:
            args.ofa_checkpoint_path += "/kernel_depth2kernel_depth_expand/phase1/checkpoint/model_best.pth.tar"
        else:  # unused
            args.ofa_checkpoint_path += "/kernel_depth2kernel_depth_expand/phase2/checkpoint/model_best.pth.tar"

        print("Start elastic expand training")

        train_elastic_expand(train, run_manager, args, validate_func_dict)

    elif args.task == "width":
        from once_for_all.elastic_nn.training.progressive_shrinking import (
            train_elastic_width_mult,
        )

        if args.phase == 1:
            args.ofa_checkpoint_path += "/kernel2kernel_depth_expand/phase2/checkpoint/model_best.pth.tar"
        elif args.phase == 2:
            args.ofa_checkpoint_path += "/kernel_depth2kernel_depth_expand_width/phase1/checkpoint/model_best.pth.tar"
        else:
            args.ofa_checkpoint_path += "/kernel_depth2kernel_depth_expand_width/phase2/checkpoint/model_best.pth.tar"

        print("Start elastic width training")

        train_elastic_width_mult(train, run_manager, args, validate_func_dict)
    else:
        raise NotImplementedError

    print("Testing all resolutions and same params networks:")
    validate_func_dict = {
        "ft_extr_type": args.ft_extr_type,
        "ft_extr_params_list": args.ft_extr_params_list,
        "width_mult_list": net.width_mult_list,
        "ks_list": net.ks_list,
        "depth_list": net.depth_list,
        "expand_list": net.expand_ratio_list,
    }
    print("Validation dict: ", validate_func_dict)
    validate(run_manager, is_test=True)
    end = time.time()
    print("Total run time: ", end - start)