def get_model(name, input_size=None, output=None):
name = name.lower()
if name == 'lenet-300-100':
model = LeNet_300_100(input_size, output)
elif name == 'lenet-5':
model = LeNet(input_size, output)
elif 'vgg' in name:
# if 'bn' in name:
if name == 'vgg11':
model = vgg11(pretrained=False, num_classes=output)
elif name == 'vgg16':
model = vgg16(pretrained=False, num_classes=output)
else:
assert False
for n, m in model.named_modules():
if hasattr(m, 'bias') and not isinstance(m, _BatchNorm):
if m.bias is not None:
if m.bias.sum() == 0:
m.bias = None
elif 'alexnet' in name:
model = AlexNet(num_classes=output)
for n, m in model.named_modules():
if hasattr(m, 'bias') and not isinstance(m, _BatchNorm):
if m.bias is not None:
if m.bias.sum() == 0:
m.bias = None
elif 'resnet' in name:
if name == 'resnet20':
model = resnet20(num_classes=output)
elif name == 'resnet32':
model = resnet32(num_classes=output)
else:
assert False
for n, m in model.named_modules():
if hasattr(m, 'bias') and not isinstance(m, _BatchNorm):
if m.bias is not None:
if m.bias.sum() == 0:
m.bias = None
else:
assert False
return model
import models
def get_param_size(model):
params = 0
for p in model.parameters():
tmp = 1
for x in p.size():
tmp *= x
params += tmp
return params
print("revnet38: {}".format(get_param_size(models.revnet38())))
print("resnet32: {}".format(get_param_size(models.resnet32())))
if args.arch == 'resnet101':
net = resnet_imagenet.resnet101()
elif args.arch == 'resnet50':
net = resnet_imagenet.resnet50()
elif args.arch == 'resnet34':
net = resnet_imagenet.resnet34()
elif args.arch == 'resnet18':
net = resnet_imagenet.resnet18()
else:
if args.arch == 'resnet110':
net = models.resnet110(num_classes=10)
elif args.arch == 'resnet56':
net = models.resnet56(num_classes=10)
elif args.arch == 'resnet32':
net = models.resnet32(num_classes=10)
elif args.arch == 'resnet20':
net = models.resnet20(num_classes=10)
if args.dataset == 'imagenet':
if args.arch == 'resnet101':
state_dict = torch.load(
os.path.join(args.pretrain_path, 'resnet101-5d3b4d8f.pth'))
elif args.arch == 'resnet50':
state_dict = torch.load(
os.path.join(args.pretrain_path, 'resnet50-19c8e357.pth'))
elif args.arch == 'resnet34':
state_dict = torch.load(
os.path.join(args.pretrain_path, 'resnet34-333f7ec4.pth'))
elif args.arch == 'resnet18':
def init_params(target):
if target == 'MNIST':
batch_size = 128
l_inf_bound = .3 if L_INF_BOUND == 'Auto' else L_INF_BOUND
n_labels = 10
n_channels = 1
target_model = models.MNIST_target_net().to(device)
train_dataset = torchvision.datasets.MNIST(
'./datasets',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = torchvision.datasets.MNIST(
'./datasets',
train=False,
transform=transforms.ToTensor(),
download=True)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
test_dataloader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
elif target == 'CIFAR10':
batch_size = 400
l_inf_bound = 8 / 255 if L_INF_BOUND == 'Auto' else L_INF_BOUND / 255
n_labels = 10
n_channels = 3
target_model = models.resnet32().to(device)
train_dataset = torchvision.datasets.CIFAR10(
'./datasets',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = torchvision.datasets.CIFAR10(
'./datasets',
train=False,
transform=transforms.ToTensor(),
download=True)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
test_dataloader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
elif target == 'HighResolution':
batch_size = 70
l_inf_bound = .01 if L_INF_BOUND == 'Auto' else L_INF_BOUND
n_labels = 1000
n_channels = 3
target_model = m.inception_v3(pretrained=True).to(device)
target_model.eval()
transform = transforms.Compose([
transforms.Resize(299),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
dataset = cd.HighResolutionDataset('./datasets/high_resolution/img',
transform=transform)
train_dataset, test_dataset = cd.split_dataset(dataset)
train_sampler = SubsetRandomSampler(train_dataset)
test_sampler = SubsetRandomSampler(test_dataset)
train_dataloader = DataLoader(dataset,
batch_size=batch_size,
sampler=train_sampler)
test_dataloader = DataLoader(dataset,
batch_size=batch_size,
sampler=test_sampler)
else:
raise NotImplementedError('Unknown Dataset')
return train_dataloader, test_dataloader, target_model, batch_size, l_inf_bound, n_labels, n_channels, len(
test_dataset)
def main():
# Init logger
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path)
if args.resume:
if not os.path.isdir(args.resume):
os.makedirs(args.resume)
log = open(os.path.join(args.save_path, '{}.txt'.format(args.description)), 'w')
print_log('save path : {}'.format(args.save_path), log)
state = {k: v for k, v in args._get_kwargs()}
print_log(state, log)
print_log("Random Seed: {}".format(args.manualSeed), log)
print_log("use cuda: {}".format(args.use_cuda), log)
print_log("python version : {}".format(sys.version.replace('\n', ' ')), log)
print_log("torch version : {}".format(torch.__version__), log)
print_log("cudnn version : {}".format(torch.backends.cudnn.version()), log)
print_log("Compress Rate: {}".format(args.rate), log)
print_log("Epoch prune: {}".format(args.epoch_prune), log)
print_log("description: {}".format(args.description), log)
# Init data loader
if args.dataset=='cifar10':
train_loader=dataset.cifar10DataLoader(True,args.batch_size,True,args.workers)
test_loader=dataset.cifar10DataLoader(False,args.batch_size,False,args.workers)
num_classes=10
elif args.dataset=='cifar100':
train_loader=dataset.cifar100DataLoader(True,args.batch_size,True,args.workers)
test_loader=dataset.cifar100DataLoader(False,args.batch_size,False,args.workers)
num_classes=100
elif args.dataset=='imagenet':
assert False,'Do not finish imagenet code'
else:
assert False,'Do not support dataset : {}'.format(args.dataset)
# Init model
if args.arch=='cifarvgg16':
net=models.vgg16_cifar(True,num_classes)
elif args.arch=='resnet32':
net=models.resnet32(num_classes)
elif args.arch=='resnet56':
net=models.resnet56(num_classes)
elif args.arch=='resnet110':
net=models.resnet110(num_classes)
else:
assert False,'Not finished'
print_log("=> network:\n {}".format(net),log)
net = torch.nn.DataParallel(net, device_ids=list(range(args.ngpu)))
# define loss function (criterion) and optimizer
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(net.parameters(), state['learning_rate'], momentum=state['momentum'],
weight_decay=state['decay'], nesterov=True)
if args.use_cuda:
net.cuda()
criterion.cuda()
recorder = RecorderMeter(args.epochs)
# optionally resume from a checkpoint
if args.resume:
if os.path.isfile(args.resume+'checkpoint.pth.tar'):
print_log("=> loading checkpoint '{}'".format(args.resume+'checkpoint.pth.tar'), log)
checkpoint = torch.load(args.resume+'checkpoint.pth.tar')
recorder = checkpoint['recorder']
args.start_epoch = checkpoint['epoch']
if args.use_state_dict:
net.load_state_dict(checkpoint['state_dict'])
else:
net = checkpoint['state_dict']
optimizer.load_state_dict(checkpoint['optimizer'])
print_log("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']), log)
if args.evaluate:
time1=time.time()
validate(test_loader,net,criterion,args.use_cuda,log)
time2=time.time()
print('validate function took %0.3f ms' % ((time2 - time1) * 1000.0))
return
else:
print_log("=> no checkpoint found at '{}'".format(args.resume), log)
else:
print_log("=> not use any checkpoint for {} model".format(args.description), log)
if args.original_train:
original_train.args.arch=args.arch
original_train.args.dataset=args.dataset
original_train.main()
return
comp_rate=args.rate
m=mask.Mask(net,args.use_cuda)
print("-" * 10 + "one epoch begin" + "-" * 10)
print("the compression rate now is %f" % comp_rate)
val_acc_1, val_los_1 = validate(test_loader, net, criterion, args.use_cuda,log)
print(" accu before is: %.3f %%" % val_acc_1)
m.model=net
print('before pruning')
m.init_mask(comp_rate,args.last_index)
m.do_mask()
print('after pruning')
m.print_weights_zero()
net=m.model#update net
if args.use_cuda:
net=net.cuda()
val_acc_2, val_los_2 = validate(test_loader, net, criterion, args.use_cuda,log)
print(" accu after is: %.3f %%" % val_acc_2)
#
start_time=time.time()
epoch_time=AverageMeter()
for epoch in range(args.start_epoch,args.epochs):
current_learning_rate=adjust_learning_rate(args.learning_rate,optimizer,epoch,args.gammas,args.schedule)
need_hour, need_mins, need_secs = convert_secs2time(epoch_time.avg * (args.epochs - epoch))
need_time = '[Need: {:02d}:{:02d}:{:02d}]'.format(need_hour, need_mins, need_secs)
print_log(
'\n==>>{:s} [Epoch={:03d}/{:03d}] {:s} [learning_rate={:6.4f}]'.format(time_string(), epoch, args.epochs,
need_time, current_learning_rate) \
+ ' [Best : Accuracy={:.2f}]'.format(recorder.max_accuracy(False)), log)
train_acc,train_los=train(train_loader,net,criterion,optimizer,epoch,args.use_cuda,log)
validate(test_loader, net, criterion,args.use_cuda, log)
if (epoch % args.epoch_prune == 0 or epoch == args.epochs - 1):
m.model=net
print('before pruning')
m.print_weights_zero()
m.init_mask(comp_rate,args.last_index)
m.do_mask()
print('after pruning')
m.print_weights_zero()
net=m.model
if args.use_cuda:
net=net.cuda()
val_acc_2, val_los_2 = validate(test_loader, net, criterion,args.use_cuda,log)
is_best = recorder.update(epoch, train_los, train_acc, val_los_2, val_acc_2)
if args.resume:
save_checkpoint({
'epoch': epoch + 1,
'state_dict': net,
'recorder': recorder,
'optimizer': optimizer.state_dict(),
}, is_best, args.resume, 'checkpoint.pth.tar')
print('save ckpt done')
epoch_time.update(time.time()-start_time)
start_time=time.time()
torch.save(net,args.model_save)
# torch.save(net,args.save_path)
flops.print_model_param_nums(net)
flops.count_model_param_flops(net,32,False)
log.close()
def main():
global args
args = parser.parse_args()
with open(args.config) as f:
config = yaml.load(f)
for key in config:
for k, v in config[key].items():
setattr(args, k, v)
transform_train = transforms.Compose([
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
# Normalize the test set same as training set without augmentation
transform_test = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])
if args.imbalance == "regular":
train_dataset = torchvision.datasets.CIFAR10(
root='./data', train=True, download=True, transform=transform_train)
else:
train_dataset = IMBALANCECIFAR10(root='../part1-convnet/data',
transform=transform_train,
)
cls_num_list = train_dataset.get_cls_num_list()
if args.reweight:
per_cls_weights = reweight(cls_num_list, beta=args.beta)
if torch.cuda.is_available():
per_cls_weights = per_cls_weights.cuda()
else:
per_cls_weights = None
train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True)
test_dataset = torchvision.datasets.CIFAR10(
root='./data', train=False, download=True, transform=transform_test)
test_loader = torch.utils.data.DataLoader(
test_dataset, batch_size=100, shuffle=False, num_workers=2)
if args.model == 'TwoLayerNet':
model = TwoLayerNet(3072, 256, 10)
elif args.model == 'VanillaCNN':
model = VanillaCNN()
elif args.model == 'MyModel':
model = MyModel()
elif args.model == 'ResNet-32':
model = resnet32()
print(model)
if torch.cuda.is_available():
model = model.cuda()
if args.loss_type == "CE":
criterion = nn.CrossEntropyLoss()
else:
criterion = FocalLoss(weight=per_cls_weights, gamma=1)
optimizer = torch.optim.SGD(model.parameters(), args.learning_rate,
momentum=args.momentum,
weight_decay=args.reg)
best = 0.0
best_cm = None
best_model = None
for epoch in range(args.epochs):
adjust_learning_rate(optimizer, epoch, args)
# train loop
train(epoch, train_loader, model, optimizer, criterion)
# validation loop
acc, cm = validate(epoch, test_loader, model, criterion)
if acc > best:
best = acc
best_cm = cm
best_model = copy.deepcopy(model)
print('Best Prec @1 Acccuracy: {:.4f}'.format(best))
per_cls_acc = best_cm.diag().detach().numpy().tolist()
for i, acc_i in enumerate(per_cls_acc):
print("Accuracy of Class {}: {:.4f}".format(i, acc_i))
if args.save_best:
torch.save(best_model.state_dict(), './checkpoints/' + args.model.lower() + '.pth')
np.random.seed(args.seed)
print('Preparing directory {}'.format(args.dir))
os.makedirs(args.dir, exist_ok=True)
with open(os.path.join(args.dir, 'command.sh'), 'w') as f:
f.write(' '.join(sys.argv))
f.write('\n')
## dataset
print('Loading dataset')
loaders = dataLoader.load_cifar10(args.batch_size)
## model
print('Building model')
num_classes = 6
model = models.resnet32()
if args.se:
model = models.se_resnet32()
use_gpu = torch.cuda.is_available()
if use_gpu:
print('Let us use {} GPUs'.format(torch.cuda.device_count()))
model = nn.DataParallel(model)
model.cuda()
## loss func
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),
lr=args.lr_init,
momentum=args.momentum,
weight_decay=args.wd)
def init_params(target):
if target == 'MNIST':
batch_size = 128
l_inf_bound = .3 if L_INF_BOUND == 'Auto' else L_INF_BOUND
n_labels = 10
n_channels = 1
target_model = models.MNIST_target_net().to(device)
train_dataset = torchvision.datasets.MNIST(
'./datasets',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = torchvision.datasets.MNIST(
'./datasets',
train=False,
transform=transforms.ToTensor(),
download=True)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
test_dataloader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
elif target == 'CIFAR10':
batch_size = 400
l_inf_bound = 8 / 255 if L_INF_BOUND == 'Auto' else L_INF_BOUND / 255
n_labels = 10
n_channels = 3
target_model = models.resnet32().to(device)
train_dataset = torchvision.datasets.CIFAR10(
'./datasets',
train=True,
transform=transforms.ToTensor(),
download=True)
test_dataset = torchvision.datasets.CIFAR10(
'./datasets',
train=False,
transform=transforms.ToTensor(),
download=True)
train_dataloader = DataLoader(train_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
test_dataloader = DataLoader(test_dataset,
batch_size=batch_size,
shuffle=True,
num_workers=1)
else:
raise NotImplementedError('Unknown Dataset')
return train_dataloader, test_dataloader, target_model, batch_size, l_inf_bound, n_labels, n_channels, len(
test_dataset)