Dubhe/dubhe-tadl/classic_nas/retrainer.py

# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import copy
import logging
import os
import argparse
import logging
import sys
sys.path.append('..'+ '/' + '..')
from collections import OrderedDict


import torch
import torch.nn as nn
import torch.nn.functional as F

from torchvision import datasets, transforms
from model import Net

from pytorch.trainer import Trainer
from pytorch.utils import AverageMeterGroup
from pytorch.utils import mkdirs
from pytorch.mutables import LayerChoice, InputChoice
from fixed import apply_fixed_architecture

from mutator import ClassicMutator
from abc import ABC, abstractmethod
from pytorch.retrainer import Retrainer
import numpy as np
import time
import json

logger = logging.getLogger(__name__)
#logger.setLevel(logging.INFO)


class ClassicnasRetrainer(Retrainer):
    """
    Classicnas trainer.

    Parameters
    ----------
    model : nn.Module
        PyTorch model to be trained.
    loss : callable
        Receives logits and ground truth label, return a loss tensor.
    metrics : callable
        Receives logits and ground truth label, return a dict of metrics.
    optimizer : Optimizer
        The optimizer used for optimizing the model.
    epochs : int
        Number of epochs planned for training.
    dataset_train : Dataset
        Dataset for training. Will be split for training weights and architecture weights.
    dataset_valid : Dataset
        Dataset for testing.
    mutator : ClassicMutator
        Use in case of customizing your own ClassicMutator. By default will instantiate a ClassicMutator.
    batch_size : int
        Batch size.
    workers : int
        Workers for data loading.
    device : torch.device
        ``torch.device("cpu")`` or ``torch.device("cuda")``.
    log_frequency : int
        Step count per logging.
    callbacks : list of Callback
        list of callbacks to trigger at events.
    arc_learning_rate : float
        Learning rate of architecture parameters.
    unrolled : float
        ``True`` if using second order optimization, else first order optimization.
    """
    def __init__(self, model, loss, metrics,
                 optimizer, epochs, dataset_train, dataset_valid, search_space_path,selected_space_path,checkpoint_dir,trial_id,
                 mutator=None, batch_size=64, workers=4, device=None, log_frequency=None,
                 callbacks=None, arc_learning_rate=3.0E-4, unrolled=False):


        self.model = model

        self.loss = loss
        self.metrics = metrics 
        self.optimizer = optimizer
        self.epochs = epochs
        self.device = device
        self.batch_size = batch_size
        self.checkpoint_dir =checkpoint_dir

        self.train_loader = torch.utils.data.DataLoader(
            datasets.MNIST(dataset_train, train=True, download=True,
                        transform=transforms.Compose([
                            transforms.ToTensor(),
                            transforms.Normalize((0.1307,), (0.3081,))
                        ])),
            batch_size=batch_size, shuffle=True, **kwargs)

        self.test_loader = torch.utils.data.DataLoader(
            datasets.MNIST(dataset_valid, train=False, transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.1307,), (0.3081,))
            ])),
            batch_size=1000, shuffle=True, **kwargs)



        self.search_space_path = search_space_path
        self.selected_space_path =selected_space_path
        self.trial_id = trial_id
        self.result = {"accuracy": [],"cost_time": 0.}

    def train(self):
        

        # t1 = time()
        # phase 1. architecture step
        #print(self.model.state_dict)
        apply_fixed_architecture(self.model, self.selected_space_path)
        #print(self.model.state_dict)
        # phase 2: child network step
        for child_epoch in range(1, self.epochs + 1):

            self.model.train()
            for batch_idx, (data, target) in enumerate(self.train_loader):
                data, target = data.to(self.device), target.to(self.device)
                optimizer.zero_grad()
                output = self.model(data)
                loss = F.nll_loss(output, target)
                loss.backward()
                optimizer.step()
                if batch_idx % args['log_interval'] == 0:
                    logger.info('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        child_epoch, batch_idx * len(data), len(self.train_loader.dataset),
                        100. * batch_idx / len(self.train_loader), loss.item()))

            test_acc = self.validate()
            print({"type":"accuracy","result":{"sequence":child_epoch,"category":"epoch","value":test_acc}} )
            with open(args['result_path'], "a") as ss_file:
                ss_file.write(json.dumps({"type":"accuracy","result":{"sequence":child_epoch,"category":"epoch","value":test_acc}} ) + '\n')
            self.result['accuracy'].append(test_acc)
    

    def validate(self):
        self.model.eval()
        test_loss = 0
        correct = 0
        with torch.no_grad():
            for data, target in self.test_loader:
                data, target = data.to(self.device), target.to(self.device)
                output = self.model(data)
                # sum up batch loss
                test_loss += F.nll_loss(output, target, reduction='sum').item()
                # get the index of the max log-probability
                pred = output.argmax(dim=1, keepdim=True)
                correct += pred.eq(target.view_as(pred)).sum().item()

        test_loss /= len(self.test_loader.dataset)

        accuracy = 100. * correct / len(self.test_loader.dataset)

        logger.info('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
            test_loss, correct, len(self.test_loader.dataset), accuracy))

        return accuracy

    # 
    # def export(self, file):
    #     """
    #     Override the method to export to file.

    #     Parameters
    #     ----------
    #     file : str
    #         File path to export to.
    #     """
    #     raise NotImplementedError

     
    def checkpoint(self):
        """
        Override to dump a checkpoint.
        """
        if isinstance(self.model, nn.DataParallel):
            state_dict = self.model.module.state_dict()
        else:
            state_dict = self.model.state_dict()
        if not os.path.exists(self.checkpoint_dir):
            os.makedirs(self.checkpoint_dir)
        dest_path = os.path.join(self.checkpoint_dir, f"best_checkpoint_epoch{self.epochs}.pth")
        logger.info("Saving model to %s", dest_path)
        torch.save(state_dict, dest_path)



def dump_global_result(args,global_result):
    with open(args['result_path'], "w") as ss_file:
        json.dump(global_result, ss_file, sort_keys=True, indent=2)

def get_params():
    # Training settings
    parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
    parser.add_argument("--data_dir", type=str,
                        default='./data', help="data directory")
    parser.add_argument("--best_selected_space_path", type=str,
                        default='./selected_space.json', help="selected_space_path")
    parser.add_argument("--search_space_path", type=str,
                        default='./selected_space.json', help="search_space_path")        
    parser.add_argument("--result_path", type=str,
                        default='./result.json', help="result_path")
    parser.add_argument('--batch_size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument("--hidden_size", type=int, default=512, metavar='N',
                        help='hidden layer size (default: 512)')
    parser.add_argument('--lr', type=float, default=0.01, metavar='LR',
                        help='learning rate (default: 0.01)')
    parser.add_argument('--momentum', type=float, default=0.5, metavar='M',
                        help='SGD momentum (default: 0.5)')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='number of epochs to train (default: 10)')
    parser.add_argument('--seed', type=int, default=1, metavar='S',
                        help='random seed (default: 1)')
    parser.add_argument('--no_cuda', default=False,
                        help='disables CUDA training')
    parser.add_argument('--log_interval', type=int, default=1000, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument("--best_checkpoint_dir",type=str,default="path/to/",
                        help="Path for saved checkpoints. (default: %(default)s)")
    parser.add_argument('--trial_id', type=int, default=0, metavar='N',
                    help='trial_id,start from 0')

    args, _ = parser.parse_known_args()
    return args

if __name__ == '__main__':
    try:
        start=time.time()

        params = vars(get_params())
        args =params

        use_cuda = not args['no_cuda'] and torch.cuda.is_available()

        torch.manual_seed(args['seed'])

        device = torch.device("cuda" if use_cuda else "cpu")

        kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}

        data_dir = args['data_dir']

        hidden_size = args['hidden_size']

        model = Net(hidden_size=hidden_size).to(device)

        optimizer = torch.optim.SGD(model.parameters(), lr=args['lr'],
                            momentum=args['momentum'])

        #mkdirs(args['search_space_path'])
        mkdirs(args['best_selected_space_path'])
        mkdirs(args['result_path'])
        trainer = ClassicnasRetrainer(model,
                        loss=None,
                        metrics=None,
                        optimizer=optimizer,
                        epochs=args['epochs'],
                        dataset_train=data_dir,
                        dataset_valid=data_dir,
                        search_space_path = args['search_space_path'],
                        selected_space_path = args['best_selected_space_path'],
                        checkpoint_dir = args['best_checkpoint_dir'],
                        trial_id = args['trial_id'],
                        batch_size=args['batch_size'],
                        log_frequency=args['log_interval'],
                        device= device,
                        unrolled=None,
                        callbacks=None)

        with open(args['result_path'], "w") as ss_file:
            ss_file.write('')
        trainer.train()
        trainer.checkpoint()
        global_result = trainer.result
        global_result['cost_time'] = str(time.time() - start) +'s'
        #dump_global_result(params,global_result)
    except Exception as exception:
        logger.exception(exception)
        raise