def main():
# args
args = getArgs()
# logging
projectName = "{}_{}_{}_{}_{}_{}".format(args.model.lower(), args.datasets,
args.epochs, args.batch_size,
args.lr, args.postfix)
modelDir = os.path.join(args.save_dir, projectName)
logger = get_logger(modelDir)
with open(os.path.join(modelDir, "args.yaml"),
"w") as yaml_file: # dump experiment config
yaml.dump(args, yaml_file)
pymonitor = ProgressMonitor(logger)
tbmonitor = TensorBoardMonitor(logger, modelDir)
monitors = [pymonitor, tbmonitor]
# dataloader
trainLoader = getDataloader(args.datasets, "train", args.batch_size,
args.workers, args.pin_memory, args.cutout)
valLoader = getDataloader(args.datasets, "val", args.batch_size,
args.workers, args.pin_memory, args.cutout)
# device init
if args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
np.random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
args.use_cuda = args.gpus > 0 and torch.cuda.is_available()
args.device = torch.device('cuda' if args.use_cuda else 'cpu')
if args.use_cuda:
torch.cuda.manual_seed(args.manualSeed)
cudnn.benchmark = True
# model
# model = vgg.__dict__[args.model](pretrained=args.pretrained,
# progress=True,
# num_classes=args.class_num,
# binarynum=args.binarynum)
model = model_dict[args.model](num_classes=args.class_num,
binarynum=args.binarynum)
logger.info("model is:{} \n".format(model))
if torch.cuda.device_count() > 1 and args.use_cuda:
logger.info('use: %d gpus', torch.cuda.device_count())
model = nn.DataParallel(model)
# loss and optimazer
criterion = nn.CrossEntropyLoss()
if args.use_cuda:
logger.info("load model and criterion to gpu !")
model = model.to(args.device)
criterion = criterion.to(args.device)
if args.optimizer.lower() == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer.lower() == 'adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
weight_decay=args.weight_decay)
else:
NotImplementedError()
if args.scheduler.lower() == "cos":
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
args.epochs,
eta_min=0,
last_epoch=-1)
elif args.scheduler.lower() == "steplr":
scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer,
step_size=args.step_size,
gamma=0.1,
last_epoch=-1)
else:
NotImplementedError()
# best recoder
perf_scoreboard = PerformanceScoreboard(args.num_best_scores)
# resume
start_epoch = 0
assert os.path.isfile(args.resume_path)
model, _, _ = loadCheckpoint(args.resume_path, model, args)
# just eval model
if args.eval:
validate(valLoader, model, criterion, -1, monitors, args, logger)
else:
# resume training or pretrained model, we should eval model firstly.
if args.resume or args.pretrained:
logger.info('>>>>>>>> Epoch -1 (pre-trained model evaluation)')
top1, top5, _ = validate(valLoader, model, criterion,
start_epoch - 1, monitors, args, logger)
logger.info(
' Model evaluation score==> Epoch [%d][Top1: %.3f Top5: %.3f]',
-1, top1, top5)
# start training
for epoch in range(start_epoch, args.epochs):
logger.info('>>>> Epoch {} Lr {}'.format(
epoch, optimizer.param_groups[0]['lr']))
t_top1, t_top5, t_loss = train(trainLoader, model, criterion,
optimizer, scheduler, epoch,
monitors, args, logger)
v_top1, v_top5, v_loss = validate(valLoader, model, criterion,
epoch, monitors, args, logger)
tbmonitor.writer.add_scalars('Train_vs_Validation/Loss', {
'train': t_loss,
'val': v_loss
}, epoch)
tbmonitor.writer.add_scalars('Train_vs_Validation/Top1', {
'train': t_top1,
'val': v_top1
}, epoch)
tbmonitor.writer.add_scalars('Train_vs_Validation/Top5', {
'train': t_top5,
'val': v_top5
}, epoch)
# add params
for name, param in model.named_parameters():
if ("conv" in name.lower()
or "fc" in name.lower()) and "weight" in name.lower():
tbmonitor.writer.add_histogram(name, param.data.cpu(),
epoch)
l, board = perf_scoreboard.update(v_top1, v_top5, epoch)
for idx in range(l):
score = board[idx]
logger.info(
'Scoreboard best %d ==> Epoch [%d][Top1: %.3f Top5: %.3f]',
idx + 1, score['epoch'], score['top1'], score['top5'])
is_best = perf_scoreboard.is_best(epoch)
# save model
if epoch % args.save_freq == 0:
saveCheckpoint(
epoch, args.model, model, {
'scheduler': scheduler,
'optimizer': optimizer,
'perf_scoreboard': perf_scoreboard
}, is_best, os.path.join(modelDir, "ckpts"))
# update lr
scheduler.step()
logger.info('>>>>>>>> Epoch -1 (final model evaluation)')
validate(valLoader, model, criterion, -1, monitors, args, logger)
tbmonitor.writer.close() # close the TensorBoard
logger.info('Program completed successfully ... exiting ...')
logger.info(
'If you have any questions or suggestions, please visit: github.com/lswzjuer/nas-pruning-quantize'
)
def main():
# args
args = getArgs()
# logging
projectName = "{}_{}_{}_{}_{}_{}".format(args.model.lower(), args.datasets,
args.epochs, args.batch_size,
args.lr, args.postfix)
modelDir = os.path.join(args.save_dir, projectName)
logger = get_logger(modelDir)
with open(os.path.join(modelDir, "args.yaml"),
"w") as yaml_file: # dump experiment config
yaml.dump(args, yaml_file)
pymonitor = ProgressMonitor(logger)
tbmonitor = TensorBoardMonitor(logger, modelDir)
monitors = [pymonitor, tbmonitor]
# dataloader
trainLoader = getDataloader(args.datasets, "train", args.batch_size,
args.workers, args.pin_memory, args.cutout)
valLoader = getDataloader(args.datasets, "val", args.batch_size,
args.workers, args.pin_memory, args.cutout)
# device init
if args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
np.random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
args.use_cuda = args.gpus > 0 and torch.cuda.is_available()
args.device = torch.device('cuda' if args.use_cuda else 'cpu')
if args.use_cuda:
torch.cuda.manual_seed(args.manualSeed)
cudnn.benchmark = True
# model
model = vgg.__dict__[args.model](pretrained=args.pretrained,
progress=True,
num_classes=args.class_num)
logger.info("model is:{} \n".format(model))
if torch.cuda.device_count() > 1 and args.use_cuda:
logger.info('use: %d gpus', torch.cuda.device_count())
model = nn.DataParallel(model)
# loss and optimazer
criterion = nn.CrossEntropyLoss()
if args.use_cuda:
logger.info("load model and criterion to gpu !")
model = model.to(args.device)
criterion = criterion.to(args.device)
# all_parameters = model.parameters()
# act_params=[]
# for pname,p in model.named_parameters():
# if "nonlinear" in pname:
# act_params.append(p)
# act_params_id= list(map(id,act_params))
# other_parameters= list(filter(lambda p: id(p) not in act_params_id, all_parameters))
logger.info("Act layer params :\n {}".format(model.alpha_parameters()))
if args.optimizer.lower() == 'sgd':
optimizer = torch.optim.SGD([{
"params": model.other_parameters(),
"lr": args.lr,
"weight_decay": args.weight_decay
}, {
"params": model.alpha_parameters(),
"lr": 0.001,
"weight_decay": 0.0
}],
momentum=args.momentum)
elif args.optimizer.lower() == 'adam':
optimizer = torch.optim.Adam([{
"params": model.other_parameters(),
"lr": args.lr,
"weight_decay": args.weight_decay
}, {
"params": model.alpha_parameters(),
"lr": 0.001,
"weight_decay": 0.0
}],
betas=(0.9, 0.999))
else:
NotImplementedError()
if args.scheduler.lower() == "cos":
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
args.epochs,
eta_min=0,
last_epoch=-1)
elif args.scheduler.lower() == "steplr":
scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer,
step_size=args.step_size,
gamma=0.1,
last_epoch=-1)
else:
NotImplementedError()
# best recoder
perf_scoreboard = PerformanceScoreboard(args.num_best_scores)
# resume
start_epoch = 0
if args.resume:
if os.path.isfile(args.resume_path):
model, extras, start_epoch = loadCheckpoint(
args.resume_path, model, args)
optimizer, scheduler, perf_scoreboard = extras[
"optimizer"], extras['scheduler'], extras["perf_scoreboard"]
else:
raise FileNotFoundError("No checkpoint found at '{}'".format(
args.resume))
T_min, T_max = Tdict[args.model]
logging.info("Tmin:{} Tmax:{}".format(T_min, T_max))
# update k,t
def Log_UP(K_min, K_max, epoch):
Kmin, Kmax = math.log(K_min) / math.log(10), math.log(
K_max) / math.log(10)
return torch.tensor(
[math.pow(10,
Kmin + (Kmax - Kmin) / args.epochs * epoch)]).float()
# just eval model
if args.eval:
validate(valLoader, model, criterion, -1, monitors, args, logger)
else:
# resume training or pretrained model, we should eval model firstly.
if args.resume or args.pretrained:
logger.info('>>>>>>>> Epoch -1 (pre-trained model evaluation)')
top1, top5, _ = validate(valLoader, model, criterion,
start_epoch - 1, monitors, args, logger)
l, board = perf_scoreboard.update(top1, top5, start_epoch - 1)
for idx in range(l):
score = board[idx]
logger.info(
'Scoreboard best %d ==> Epoch [%d][Top1: %.3f Top5: %.3f]',
idx + 1, score['epoch'], score['top1'], score['top5'])
# start training
for epoch in range(start_epoch, args.epochs):
t = Log_UP(T_min, T_max, epoch)
if (t < 1):
k = 1 / t
else:
k = torch.tensor([1]).float()
if args.use_cuda:
k = k.cuda()
t = t.cuda()
model.conv0.k = k
model.conv1.k = k
model.conv2.k = k
model.conv3.k = k
model.conv4.k = k
model.conv5.k = k
model.conv0.t = t
model.conv1.t = t
model.conv2.t = t
model.conv3.t = t
model.conv4.t = t
model.conv5.t = t
logger.info('>>>> Epoch {} Lr {} K:{} T:{}'.format(
epoch, optimizer.param_groups[0]['lr'], k, t))
t_top1, t_top5, t_loss = train(trainLoader, model, criterion,
optimizer, scheduler, epoch,
monitors, args, logger)
v_top1, v_top5, v_loss = validate(valLoader, model, criterion,
epoch, monitors, args, logger)
tbmonitor.writer.add_scalars('Train_vs_Validation/Loss', {
'train': t_loss,
'val': v_loss
}, epoch)
tbmonitor.writer.add_scalars('Train_vs_Validation/Top1', {
'train': t_top1,
'val': v_top1
}, epoch)
tbmonitor.writer.add_scalars('Train_vs_Validation/Top5', {
'train': t_top5,
'val': v_top5
}, epoch)
l, board = perf_scoreboard.update(v_top1, v_top5, epoch)
for idx in range(l):
score = board[idx]
logger.info(
'Scoreboard best %d ==> Epoch [%d][Top1: %.3f Top5: %.3f]',
idx + 1, score['epoch'], score['top1'], score['top5'])
is_best = perf_scoreboard.is_best(epoch)
# save model
if epoch % args.save_freq == 0:
saveCheckpoint(
epoch, args.model, model, {
'scheduler': scheduler,
'optimizer': optimizer,
'perf_scoreboard': perf_scoreboard
}, is_best, os.path.join(modelDir, "ckpts"))
# update lr
scheduler.step()
logger.info('>>>>>>>> Epoch -1 (final model evaluation)')
validate(valLoader, model, criterion, -1, monitors, args, logger)
tbmonitor.writer.close() # close the TensorBoard
logger.info('Program completed successfully ... exiting ...')
logger.info(
'If you have any questions or suggestions, please visit: github.com/lswzjuer/nas-pruning-quantize'
)
def main():
# args
args = getArgs()
if args.arch != "resnet18":
args.steplist = [150, 220, 260]
else:
args.steplist = [40, 80, 120]
# logging
projectName = "{}_{}_{}_{}_{}_{}".format(args.model.lower(), args.datasets,
args.epochs, args.batch_size,
args.lr, args.postfix)
modelDir = os.path.join(args.save_dir, projectName)
logger = get_logger(modelDir)
with open(os.path.join(modelDir, "args.yaml"),
"w") as yaml_file: # dump experiment config
yaml.dump(args, yaml_file)
pymonitor = ProgressMonitor(logger)
tbmonitor = TensorBoardMonitor(logger, modelDir)
monitors = [pymonitor, tbmonitor]
# dataloader
trainLoader = getDataloader(args.datasets, "train", args.batch_size,
args.workers, args.pin_memory, args.cutout)
valLoader = getDataloader(args.datasets, "val", args.batch_size,
args.workers, args.pin_memory, args.cutout)
# device init
if args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
np.random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
args.device = torch.device('cuda')
# model
model = ARCH_DICT[args.arch].__dict__[args.model](
pretrained=args.pretrained, progress=True, num_classes=args.class_num)
logger.info("model is:{} \n".format(model))
# loss and optimazer
criterion = nn.CrossEntropyLoss()
# best recoder
perf_scoreboard = PerformanceScoreboard(args.num_best_scores)
# # resume
# assert args.resume == True
# args.resume_path = r"F:\source_code\nas-pruning-quantize\network_quantize\checkpoints\BNN\resnet20_cifar10_300_128_0.1_mstep\ckpts\resnet20_best.pth.tar"
# args.resume_path = r"F:\source_code\nas-pruning-quantize\network_quantize\checkpoints\BNN\vgg_small_cifar10_300_128_0.1_mstep\ckpts\vgg_small_best.pth.tar"
args.resume_path = r"F:\source_code\nas-pruning-quantize\network_quantize\checkpoints\BNN\resnet18_cifar10_150_128_0.1_mstep_d8\ckpts\resnet18_best.pth.tar"
model, extras, start_epoch = loadCheckpoint(args.resume_path, model, args)
_, _, _ = extras["optimizer"], extras['scheduler'], extras[
"perf_scoreboard"]
# # just eval model
tp1, tp5, loss = validate(valLoader, model, criterion, args)
print("origianl model eval :{} ".format(tp1, tp5, loss))
# 量化bits[8,6,5,4,3,2,1]
# 逐层进行量化
# vggsmall,resnet18,resnet20
careops = ["Conv2d", "Linear"]
testbits = [8, 7, 6, 5, 4, 3, 2, 1]
quantityOps = findQuantityOps(model, careops)
# 生成各个bit的逐层量化方案
quantityinfo = createQuantityInfo(quantityOps, testbits)
# 根据生成的量化方案执行量化分析，查看逐层量化时的准确率和loss变化情况
validateDict = quantityAnalysis(model, quantityOps, quantityinfo,
valLoader, criterion, args)
file = "./{}.pkl".format(args.model)
with open(file, "wb") as f:
pickle.dump(validateDict, f)
def main():
# args
args = getArgs()
if args.arch in ["resnet18", "mobilenetv2", "mobilenetv1", "ghostnet"]:
args.steplist = [30, 60, 90]
# args.steplist = [40,80,120]
else:
args.steplist = [150, 220, 260]
# logging
projectName = "{}_{}_{}_{}_{}_{}and{}_{}".format(
args.model.lower(), args.datasets, args.epochs, args.batch_size,
args.lr, args.wbit, args.abit, args.postfix)
modelDir = os.path.join(args.save_dir, projectName)
logger = get_logger(modelDir)
with open(os.path.join(modelDir, "args.yaml"),
"w") as yaml_file: # dump experiment config
yaml.dump(args, yaml_file)
pymonitor = ProgressMonitor(logger)
tbmonitor = TensorBoardMonitor(logger, modelDir)
monitors = [pymonitor, tbmonitor]
# dataloader
trainLoader = getDataloader(args.datasets, "train", args.batch_size,
args.workers, args.pin_memory, args.cutout)
valLoader = getDataloader(args.datasets, "val", args.batch_size,
args.workers, args.pin_memory, args.cutout)
# device init
if args.manualSeed is None:
args.manualSeed = random.randint(1, 10000)
np.random.seed(args.manualSeed)
torch.manual_seed(args.manualSeed)
args.use_cuda = args.gpus > 0 and torch.cuda.is_available()
args.device = torch.device('cuda' if args.use_cuda else 'cpu')
if args.use_cuda:
torch.cuda.manual_seed(args.manualSeed)
cudnn.benchmark = True
# model
model = ARCH_DICT[args.arch].__dict__[args.model](
pretrained=args.pretrained,
progress=True,
wbit=args.wbit,
abit=args.abit,
num_classes=args.class_num)
logger.info("model is:{} \n".format(model))
# loss and optimazer
criterion = nn.CrossEntropyLoss()
if args.use_cuda:
logger.info("load model and criterion to gpu !")
model = model.to(args.device)
criterion = criterion.to(args.device)
if args.optimizer.lower() == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
elif args.optimizer.lower() == 'adam':
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
betas=(0.9, 0.999),
weight_decay=args.weight_decay)
else:
NotImplementedError()
if args.scheduler.lower() == "cos":
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
args.epochs,
eta_min=0,
last_epoch=-1)
elif args.scheduler.lower() == "step":
scheduler = torch.optim.lr_scheduler.StepLR(optimizer=optimizer,
step_size=args.step_size,
gamma=0.1,
last_epoch=-1)
elif args.scheduler.lower() == "mstep":
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer,
args.steplist,
gamma=args.gamma)
else:
NotImplementedError()
# resume
start_epoch = 0
if args.resume:
if os.path.isfile(args.resume_path):
model, extras, start_epoch = loadCheckpoint(
args.resume_path, model, args)
_, _, _ = extras["optimizer"], extras['scheduler'], extras[
"perf_scoreboard"]
else:
raise FileNotFoundError("No checkpoint found at '{}'".format(
args.resume))
model_adv = AttackPGD(model, None)
# just eval model
if args.eval:
v1, v5, loss = validate(valLoader,
model_adv,
criterion,
args,
adv=False)
logger.info('==> Top1: %.3f Top5: %.3f Loss: %.3f\n', v1, v5,
loss)
v1, v5, loss = validate(valLoader,
model_adv,
criterion,
args,
adv=True)
logger.info('==> PGD20 Top1: %.3f Top5: %.3f Loss: %.3f\n', v1,
v5, loss)
else:
# adv training some epochs and eval
for epoch in range(args.epochs):
logger.info('>>>> Epoch {} Lr {} '.format(
epoch, optimizer.param_groups[0]['lr']))
t_top1, t_top5, t_loss = train(trainLoader,
model_adv,
criterion,
optimizer,
scheduler,
args,
adv=True)
logger.info(
'ADV Training ==> Top1: %.3f Top5: %.3f Loss: %.3f\n',
t_top1, t_top5, t_loss)
v_top1, v_top5, v_loss = validate(valLoader,
model_adv,
criterion,
args,
adv=True)
logger.info(
'ADV Validation ==> Top1: %.3f Top5: %.3f Loss: %.3f\n',
v_top1, v_top5, v_loss)
