main_train.py

import argparse
import json
import os
import logging
import time
from datetime import datetime

import torch

from load_dataset import load_dataset
from tools.utils import setup_logging, save_checkpoint, AverageMeter, accuracy
import models.spk_model



parser = argparse.ArgumentParser(description='PyTorch Training')
# args of datasets
parser.add_argument('--dataset', default='cifar10', type=str,
                    help='dataset: [cifar10|cifar100|svhn|imagenet]')
parser.add_argument('--data', default='/datasets/imagenet',
                    help='path to dataset, /datasets/imagenet,/data/public/imagenet2012')
parser.add_argument('-j', '--workers', default=16, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--auto_aug', action='store_true', default=False, help='')
parser.add_argument('--cutout', action='store_true', default=False, help='')
# args of networks
parser.add_argument('--arch', default='spk_cifar_resnet_local', type=str)
parser.add_argument('--net', default='spk_resnet18', type=str,
                    help='networks')

parser.add_argument('--seed', default=9, type=int,
                    help='seed for initializing training. ')
parser.add_argument('--save_path', default='', type=str, help='the directory used to save the trained models')
parser.add_argument('--name', default='', type=str,
                    help='name of experiment')

parser.add_argument('--epochs', default=400, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start_epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch_size', default=512, type=int,
                    metavar='N',
                    help='mini-batch size (default: 256), this is the total '
                         'batch size of all GPUs on the current node when '
                         'using Data Parallel or Distributed Data Parallel')
parser.add_argument('-p', '--print_freq', default=100, type=int,
                    metavar='N', help='print frequency (default: 10)')

parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
                    help='evaluate model on validation set')
parser.add_argument('--fine_tune', action='store_true', default=False, help='')
parser.add_argument('--save_ckpt', action='store_true', default=False, help='')

# args of optimizer
parser.add_argument('--optim', default='sgd', type=str, help='optimizer (default: sgd)')
parser.add_argument('--lr', '--learning_rate', default=0.4, type=float,
                    metavar='LR', help='initial learning rate', dest='lr')
parser.add_argument('--wd', '--weight_decay', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)',
                    dest='weight_decay')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
# Cosine learning rate
parser.add_argument('--cos_lr', dest='cos_lr', action='store_true',
                    help='whether to use cosine learning rate')
parser.set_defaults(cos_lr=False)

# args of spiking neural networks
parser.add_argument('--threshold', type=float, default=1., help='neuronal threshold (default: 1)')
parser.add_argument('--time_window', type=int, default=1, help='total time steps (default: 10)')
parser.add_argument('--decay', type=float, default=0.2, help='tau factor (default: 1)')
parser.add_argument('--pos_bn', type=int, default=3)

# args of local learning
parser.add_argument('--rule', default='AugLocal', type=str,
                    help='DECOLLE|ELL|')
parser.add_argument('--aux_net_depth', default=1, type=int,
                    help='')
parser.add_argument('--aux_net_widen', default=1.0, type=float,
                    help='widen factor of the two auxiliary nets (default: 1.0)')
parser.add_argument('--aux_net_widen_aux_loss', default=1.0, type=float,
                    help='widen factor of the two auxiliary nets (default: 1.0)')
parser.add_argument('--aux_net_feature_dim', default=128, type=int,
                    help='number of hidden features in auxiliary classifier / contrastive head '
                         '(default: 128)')
parser.add_argument('--hidden_dim', default=512, type=int,
                    help='number of hidden features in auxiliary classifier / contrastive head '
                         '(default: 128)')
parser.add_argument('--hidden_dim_aux_loss', default=512, type=int,
                    help='number of hidden features in auxiliary classifier / contrastive head '
                         '(default: 128)')
parser.add_argument('--local_module_num', default=9, type=int, help='number of local modules (1 refers to end-to-end training)')
parser.add_argument('--aux_net_config', default='unifSamp', type=str, help='architecture of auxiliary networks')
parser.add_argument('--aux_net_config_aux_loss', default='1c', type=str, help='architecture of auxiliary networks')

parser.add_argument('--detach_mem', action='store_true', default=False, help='')
parser.add_argument('--detach_reset', action='store_true', default=False, help='')
parser.add_argument('--hard_reset', action='store_true', default=False, help='')
parser.add_argument('--neuron', default='LIF', type=str, help='')




def main():
    args = parser.parse_args()
    if args.save_path == '':
        save_path = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')
        save_path = save_path + args.name + '_' + str(args.seed)
    else:
        save_path = args.save_path

    if not os.path.exists(save_path):
        os.makedirs(save_path)
    # Logging settings
    setup_logging(os.path.join(save_path, 'log.txt'))
    logging.info('saving to:' + str(save_path))
    if 'local' in args.arch:
        logging.info('local module num:' + str(args.local_module_num))


    is_cuda = torch.cuda.is_available()
    assert is_cuda, 'CPU is not supported!'
    device = torch.device('cuda' if is_cuda else 'cpu')
    from tools import set_random_seed
    set_random_seed(seed=args.seed, is_ddp=False)  # TODO
    torch.backends.cudnn.benchmark = False
    args.gpu = 'cuda'

    if args.fine_tune:
        args.weight_decay = 0.0
        # args.momentum = 0

    # Load datasets
    train_loader, val_loader, num_classes = load_dataset(args=args)
    args.num_classes = num_classes

    with open(save_path + '/args.json', 'w') as fid:
        json.dump(args.__dict__, fid, indent=2)

    logging.info('args:' + str(args))

    # Load spiking model
    model = models.spk_model.SpkModel(args=args)
    print(model)
    if args.optim == 'sgd':
        optimizer = torch.optim.SGD(model.parameters(), lr=args.lr, weight_decay=args.weight_decay,
                                    momentum=args.momentum)
    elif args.optim == 'adam':
        optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    else:
        raise NotImplementedError
    assert args.cos_lr
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, eta_min=0, T_max=args.epochs)

    best_acc1 = 0

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            # Map model to be loaded to specified single gpu.
            loc = 'cuda:{}'.format(device)
            checkpoint = torch.load(args.resume, map_location='cpu')
            logging.info(f"best acc1: {checkpoint['best_acc']}")
            state_dict = checkpoint['state_dict']
            for (key, value) in list(state_dict.items()):
                if key.startswith('module.'):
                    state_dict[key.replace("module.", "")] = value
                del state_dict[key]

            msg = model.load_state_dict(state_dict, strict=False)
            logging.info(msg)
            logging.info("=> loaded checkpoint '{}' (epoch {})"
                         .format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))

    model = torch.nn.DataParallel(model).cuda()

    criterion = torch.nn.CrossEntropyLoss()
    local_train(train_loader, val_loader, model, criterion, optimizer, scheduler, save_path, best_acc1, args)


def local_train(train_loader, val_loader, model, criterion, optimizer, scheduler, save_path, best_prec1, args):
    if args.evaluate:
        val_acc1 = local_validate_one_epoch(val_loader, model, 0, args)
        print(val_acc1)
        exit()

    for epoch in range(args.start_epoch, args.epochs):
        logging.info(f"lr: {optimizer.param_groups[0]['lr']}")
        # train for one epoch
        train_acc = local_train_one_epoch(train_loader, model, optimizer, epoch, args)
        scheduler.step()
        prec1 = local_validate_one_epoch(val_loader, model, epoch, args)
        logging.info('train acc {:.4f} test acc {:.4f}'.format(train_acc, prec1))
        # remember best prec@1 and save checkpoint
        is_best = prec1 > best_prec1
        best_prec1 = max(prec1, best_prec1)
        save_checkpoint({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'best_acc': best_prec1,
            'optimizer': optimizer.state_dict(),
        }, is_best, filename=os.path.join(save_path, 'checkpoint.pth.tar'), save_path=save_path, epoch=epoch+1)

    logging.info('Best accuracy: ' + str(best_prec1))


def local_train_one_epoch(train_loader, model, optimizer, epoch, args):
    """Train for one epoch on the training set"""
    batch_time = AverageMeter('Time', ':6.3f')
    top1 = [AverageMeter('Acc@1', ':6.2f') for _ in range(1)]
    top5 = [AverageMeter('Acc@1', ':6.2f') for _ in range(1)]
    train_batches_num = len(train_loader)
    # switch to train mode
    model.train()
    end = time.time()
    for i, (x, target) in enumerate(train_loader):
        target = target.cuda()
        if args.dataset not in ['cifar10dvs', 'hardvs']:
            x = x.cuda()
        optimizer.zero_grad()
        output = model(img=x, target=target)
        optimizer.step()
        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
        prec1, prec5 = prec1[0], prec5[0]
        top1[-1].update(prec1.item(), target.size(0))
        top5[-1].update(prec5.item(), target.size(0))
        batch_time.update(time.time() - end)
        end = time.time()
        if (i + 1) == len(train_loader) or ( i + 1 ) % args.print_freq == 0:
            string = ('Epoch: [{0}][{1}/{2}]\t'
                      'Time ({batch_time.avg:.3f})\t'
                      'Prec@1 ({top1.avg:.3f})\t'
                      'Prec@5 ({top5.avg:.3f})'.format(
                epoch, i + 1, train_batches_num, batch_time=batch_time,
                top1=top1[-1], top5=top5[-1]))
            string += (
                'mem={:.0f}MiB, max_mem={:.0f}MiB, \n'.format(torch.cuda.memory_allocated() / 1e6,
                                                                                torch.cuda.max_memory_allocated() / 1e6
                                                                                ))
            string += 'layer [{0}], Prec@1 ({top1.avg:.3f}) Prec@5 ({top5.avg:.3f}) \n'.format(0, top1=top1[-1], top5=top5[-1])
            logging.info(string)
    return top1[-1].avg


def local_validate_one_epoch(val_loader, model, epoch, args):
    """Perform validation on the validation set"""
    batch_time = AverageMeter('Time', ':6.3f')
    top1 = [AverageMeter('Acc@1', ':6.2f') for _ in range(1)]
    top5 = [AverageMeter('Acc@5', ':6.2f') for _ in range(1)]
    train_batches_num = len(val_loader)
    # switch to evaluate mode
    model.eval()
    end = time.time()
    for i, (input, target) in enumerate(val_loader):
        target = target.cuda()
        if args.dataset not in ['cifar10dvs', 'dvsgesture', 'hardvs']:
            input = input.cuda()
        with torch.no_grad():
            output = model(img=input, target=target)
        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1,5))
        prec1, prec5 = prec1[0], prec5[0]
        top1[-1].update(prec1.item(), target.size(0))
        top5[-1].update(prec5.item(), target.size(0))
        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

    string = ('Test: [{0}][{1}/{2}]\t'
              'Time ({batch_time.avg:.3f})\t'
              'Prec@1 ({top1.avg:.3f})\t'
              'Prec@5 ({top5.avg:.3f})\n'.format(
        epoch, (i + 1), train_batches_num, batch_time=batch_time,
        top1=top1[-1], top5=top5[-1]))
    # for module_num in range(1):
    #     string += 'layer [{0}], Prec@1  ({top1.avg:.3f}) Prec@5  ({top5.avg:.3f})\n'.format(module_num, top1=top1[module_num], top5=top5[module_num])
    logging.info(string)

    return top1[-1].avg


if __name__ == '__main__':
    main()