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))))