Dubhe/dubhe-tadl/pdarts/pdarts_retrain.py

import sys
sys.path.append('..'+ '/' + '..')
import os
import logging
import  time
import json
from argparse import ArgumentParser

import torch
import torch.nn as nn
# from torch.utils.tensorboard import SummaryWriter

from model import CNN
from pytorch.fixed import apply_fixed_architecture
from pytorch.utils import set_seed, mkdirs, init_logger, save_best_checkpoint, AverageMeter
from pytorch.darts import utils
from pytorch.darts import datasets
from pytorch.retrainer import Retrainer

logger = logging.getLogger(__name__)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# writer = SummaryWriter()

class PdartsRetrainer(Retrainer):
    def __init__(self, aux_weight, grad_clip, epochs, log_frequency):
        self.aux_weight = aux_weight
        self.grad_clip = grad_clip
        self.epochs = epochs
        self.log_frequency = log_frequency

    def train(self, train_loader, model, optimizer, criterion, epoch):
        top1 = AverageMeter("top1")
        top5 = AverageMeter("top5")
        losses = AverageMeter("losses")

        cur_step = epoch * len(train_loader)
        cur_lr = optimizer.param_groups[0]["lr"]
        logger.info("Epoch %d LR %.6f", epoch, cur_lr)
        # writer.add_scalar("lr", cur_lr, global_step=cur_step)

        model.train()

        for step, (x, y) in enumerate(train_loader):
            x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
            bs = x.size(0)

            optimizer.zero_grad()
            logits, aux_logits = model(x)
            loss = criterion(logits, y)
            if self.aux_weight > 0.:
                loss += self.aux_weight * criterion(aux_logits, y)
            loss.backward()
            # gradient clipping
            nn.utils.clip_grad_norm_(model.parameters(), self.grad_clip)
            optimizer.step()

            accuracy = utils.accuracy(logits, y, topk=(1, 5))
            losses.update(loss.item(), bs)
            top1.update(accuracy["acc1"], bs)
            top5.update(accuracy["acc5"], bs)
            # writer.add_scalar("loss/train", loss.item(), global_step=cur_step)
            # writer.add_scalar("acc1/train", accuracy["acc1"], global_step=cur_step)
            # writer.add_scalar("acc5/train", accuracy["acc5"], global_step=cur_step)

            if step % self.log_frequency == 0 or step == len(train_loader) - 1:
                logger.info(
                    "Train: [{:3d}/{}] Step {:03d}/{:03d} Loss {losses.avg:.3f} "
                    "Prec@(1,5) ({top1.avg:.1%}, {top5.avg:.1%})".format(
                        epoch + 1, self.epochs, step, len(train_loader) - 1, losses=losses,
                        top1=top1, top5=top5))

            cur_step += 1

        logger.info("Train: [{:3d}/{}] Final Prec@1 {:.4%}".format(epoch + 1, self.epochs, top1.avg))


    def validate(self, valid_loader, model, criterion, epoch, cur_step):
        top1 = AverageMeter("top1")
        top5 = AverageMeter("top5")
        losses = AverageMeter("losses")

        model.eval()

        with torch.no_grad():
            for step, (X, y) in enumerate(valid_loader):
                X, y = X.to(device, non_blocking=True), y.to(device, non_blocking=True)
                bs = X.size(0)

                logits = model(X)
                loss = criterion(logits, y)

                accuracy = utils.accuracy(logits, y, topk=(1, 5))
                losses.update(loss.item(), bs)
                top1.update(accuracy["acc1"], bs)
                top5.update(accuracy["acc5"], bs)

                if step % self.log_frequency == 0 or step == len(valid_loader) - 1:
                    logger.info(
                        "Valid: [{:3d}/{}] Step {:03d}/{:03d} Loss {losses.avg:.3f} "
                        "Prec@(1,5) ({top1.avg:.1%}, {top5.avg:.1%})".format(
                            epoch + 1, self.epochs, step, len(valid_loader) - 1, losses=losses,
                            top1=top1, top5=top5))

        # writer.add_scalar("loss/test", losses.avg, global_step=cur_step)
        # writer.add_scalar("acc1/test", top1.avg, global_step=cur_step)
        # writer.add_scalar("acc5/test", top5.avg, global_step=cur_step)

        logger.info("Valid: [{:3d}/{}] Final Prec@1 {:.4%}".format(epoch + 1, self.epochs, top1.avg))

        return top1.avg


if __name__ == "__main__":
    parser = ArgumentParser("Pdarts retrain")
    parser.add_argument("--data_dir", type=str,
                        default='./', help="search_space json file")
    parser.add_argument("--result_path", type=str,
                        default='./result.json', help="training result")
    parser.add_argument("--log_path", type=str,
                        default='.0/log', help="log for info")
    parser.add_argument("--best_selected_space_path", type=str,
                        default='./best_selected_space.json', help="final best selected space")
    parser.add_argument("--best_checkpoint_dir", type=str,
                        default='', help="default name is best_checkpoint_epoch{}.pth")
    parser.add_argument('--trial_id', type=int, default=0, metavar='N',
                        help='trial_id,start from 0')
    parser.add_argument("--layers", default=20, type=int)
    parser.add_argument("--batch_size", default=96, type=int)
    parser.add_argument("--log_frequency", default=10, type=int)
    parser.add_argument("--epochs", default=600, type=int)
    parser.add_argument("--lr", default=0.025, type=float)
    parser.add_argument("--channels", default=36, type=int)
    parser.add_argument("--aux_weight", default=0.4, type=float)
    parser.add_argument("--drop_path_prob", default=0.3, type=float)
    parser.add_argument("--workers", default=4)
    parser.add_argument("--grad_clip", default=5., type=float)
    args = parser.parse_args()
    
    mkdirs(args.result_path, args.log_path, args.best_checkpoint_dir)
    init_logger(args.log_path)
    logger.info(args)
    set_seed(args.trial_id) 
    logger.info("loading data")
    dataset_train, dataset_valid = datasets.get_dataset("cifar10", cutout_length=16, root=args.data_dir)

    model = CNN(32, 3, 36, 10, args.layers, auxiliary=True, search=False, dropout_rate=0.0)
    if isinstance(args.best_selected_space_path, str):
        with open(args.best_selected_space_path) as f:
            fixed_arc = json.load(f)
    apply_fixed_architecture(model, fixed_arc=fixed_arc["best_selected_space"])
    criterion = nn.CrossEntropyLoss()

    model.to(device)
    criterion.to(device)

    optimizer = torch.optim.SGD(model.parameters(), args.lr, momentum=0.9, weight_decay=3.0E-4)
    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.epochs, eta_min=1E-6)

    train_loader = torch.utils.data.DataLoader(dataset_train,
                                               batch_size=args.batch_size,
                                               shuffle=True,
                                               num_workers=args.workers,
                                               pin_memory=True)
    valid_loader = torch.utils.data.DataLoader(dataset_valid,
                                               batch_size=args.batch_size,
                                               shuffle=False,
                                               num_workers=args.workers,
                                               pin_memory=True)
    retrainer = PdartsRetrainer(aux_weight=args.aux_weight,
                                grad_clip=args.grad_clip,
                                epochs=args.epochs,
                                log_frequency = args.log_frequency)
    best_top1 = 0.
    start_time = time.time()
    for epoch in range(args.epochs):
        drop_prob = args.drop_path_prob * epoch / args.epochs
        model.drop_path_prob(drop_prob)

        # training
        retrainer.train(train_loader, model, optimizer, criterion, epoch)

        # validation
        cur_step = (epoch + 1) * len(train_loader)
        top1 = retrainer.validate(valid_loader, model, criterion, epoch, cur_step)
        # 后端在终端过滤，{"type": "Accuracy", "result": {"sequence": 1, "category": "epoch", "value":96.7}}
        logger.info({"type": "Accuracy", "result": {"sequence": epoch, "category": "epoch", "value": top1}})
        with open(args.result_path, "a") as file:
            file.write(str({"type": "Accuracy", "result": {"sequence": epoch, "category": "epoch", "value": top1}}) + '\n')
        best_top1 = max(best_top1, top1)

        lr_scheduler.step()

    logger.info("Final best Prec@1 = {:.4%}".format(best_top1))
    cost_time = time.time() - start_time
    # 后端在终端过滤，{"type": "Cost_time", "result": {"value": "* s"}}
    logger.info({"type": "Cost_time", "result": {"value": str(cost_time) + ' s'}})
    with open(args.result_path, "a") as file:
        file.write(str({"type": "Cost_time", "result": {"value": str(cost_time) + ' s'}}))
    
    # result["Cost_time"] = str(cost_time) + ' s'
    # dump_global_result(args.result_path, result)
    save_best_checkpoint(args.best_checkpoint_dir, model, optimizer, epoch)
    logger.info("Save best checkpoint in {}".format(os.path.join(args.best_checkpoint_dir, "best_checkpoint_epoch{}.pth".format(epoch))))