def networks(network, **kwargs):
# ResNet
if 'resnet' in network and 'pre' not in network:
depth = int(network[6:])
return resnet(depth, **kwargs)
elif 'vgg' in network:
depth = int(network[3:5])
if 'bn' in network:
return vgg_bn(depth, **kwargs)
else:
return vgg(depth, **kwargs)
elif 'wideResNet' in network:
depth = int(network[10:12])
widen_factor = int(network[13:])
return wideResNet(depth, widen_factor, **kwargs)
elif 'preresnet' in network:
depth = int(network[9:])
return preresnet(depth, **kwargs)
elif 'pyramidnet' in network:
depth = int(network[10:])
return pyramidnet(depth, **kwargs)
def init_network(config, num_classes):
kwargs = {'depth': config.depth, 'num_classes': num_classes}
if config.network == 'vgg':
if 'avg_pool2d' in config:
kwargs.update({'avg_pool2d': config.avg_pool2d})
if 'batch_norm' in config:
kwargs.update({'batch_norm': config.batch_norm})
net = vgg(**kwargs)
elif config.network == 'resnet':
dataset = config.data_dir.split('/')[-1]
if 'cifar' in dataset.lower():
kwargs.update({'dataset': 'cifar'})
net = resnet(**kwargs)
elif config.network == 'reskipnet':
net = reskipnet(**kwargs)
else:
raise NotImplementedError
assert os.path.exists('{}/checkpoint/original'.format(config.result_dir)),\
'No checkpoint directory for original model!'
dataset = config.data_dir.split('/')[-1]
path_to_add = '{}/{}/{}'.format(dataset, config.network, config.depth)
checkpoint_path, exp_name, epochs = get_best_checkpoint(
'{}/checkpoint/original/{}'.format(config.result_dir, path_to_add))
checkpoint = torch.load(checkpoint_path, map_location='cpu')
net.load_state_dict(checkpoint)
if torch.cuda.is_available():
net.cuda()
return net, exp_name, epochs
def init_network(args, num_classes):
kwargs = {'depth': args.depth, 'num_classes': num_classes}
depth = str(args.depth)
if args.network == 'vgg':
kwargs.update({'avg_pool2d': args.avg_pool2d})
kwargs.update({'batch_norm': True})
net = vgg(**kwargs)
elif args.network == 'resnet':
dataset = args.data_dir.split('/')[-1]
if 'cifar' in dataset.lower():
kwargs.update({'dataset': 'cifar'})
net = resnet(**kwargs)
elif args.network == 'reskipnet':
net = reskipnet(**kwargs)
elif args.network == 'resnext':
dataset = args.data_dir.split('/')[-1]
if 'cifar' in dataset.lower():
kwargs.update({'dataset': 'cifar'})
kwargs.update({
'cardinality': args.cardinality,
'base_width': args.base_width
})
net = resnext(**kwargs)
depth = '{}_{}_{}d'.format(depth, args.cardinality, args.base_width)
else:
raise NotImplementedError
if args.resume:
dataset = args.data_dir.split('/')[-1]
path_to_add = '{}/{}/{}'.format(dataset, args.network, depth)
assert os.path.exists('{}/checkpoint/original/{}'.format(args.result_dir, path_to_add)),\
'No checkpoint directory for original model!'
checkpoint_path, _, previous_epochs = get_best_checkpoint(
'{}/checkpoint/original/{}'.format(args.result_dir, path_to_add))
checkpoint = torch.load(checkpoint_path, map_location='cpu')
net.load_state_dict(checkpoint)
previous_accuracy = checkpoint_path.split('_')[-2]
else:
previous_accuracy = 0
previous_epochs = 0
if torch.cuda.is_available():
net.cuda()
return net, float(previous_accuracy), previous_epochs, depth
def __init__(self):
super(Net, self).__init__()
self.features = vgg('vgg11')
self.classifier = nn.Linear(512, 10)
net = fc(width=args.width,
depth=args.depth,
num_classes=num_classes,
input_dim=3 * 32 * 32).to(args.device)
elif args.model == 'alexnet':
if args.dataset == 'mnist':
net = alexnet(input_height=28,
input_width=28,
input_channels=1,
num_classes=num_classes)
else:
net = alexnet(ch=args.scale,
num_classes=num_classes).to(args.device)
elif args.model == 'vgg':
net = vgg(depth=args.depth,
num_classes=num_classes,
batch_norm=args.bn).to(args.device)
print(net)
opt = getattr(optim, args.optim)(net.parameters(), lr=args.lr)
if args.lr_schedule:
milestone = int(args.iterations / 3)
scheduler = optim.lr_scheduler.MultiStepLR(
opt, milestones=[milestone, 2 * milestone], gamma=0.5)
if args.criterion == 'NLL':
crit = nn.CrossEntropyLoss().to(args.device)
elif args.criterion == 'linear_hinge':
crit = linear_hinge_loss
])
dataset = datasets.ImageFolder(args.dataset_dir, transform=transform)
loader = torch.utils.data.DataLoader(dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=2)
Z_dim = 64
disc_iters = 5
discriminator = models.Discriminator().cuda()
generator = models.Generator(Z_dim * 2).cuda()
encoder_c = models.Encoder(size_num=Z_dim).cuda()
encoder_t = models.Encoder(size_num=Z_dim).cuda()
v2 = models.vgg().cuda()
optim_disc = optim.Adam(filter(lambda p: p.requires_grad,
discriminator.parameters()),
lr=args.lr,
betas=(0.5, 0.9))
optim_gen = optim.Adam(generator.parameters(), lr=args.lr, betas=(0.5, 0.999))
optim_enc_c = optim.Adam([{
'params': encoder_c.parameters()
}],
lr=args.lr,
betas=(0.9, 0.999))
optim_enc_t = optim.Adam([{
'params': encoder_t.parameters()
}],
lr=args.lr,
def main():
# Init logger
if not os.path.isdir(args.save_path):
os.makedirs(args.save_path)
log = open(
os.path.join(args.save_path, 'log_seed_{}.txt'.format(args.seed)), '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.seed), 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("Norm Pruning Rate: {}".format(args.rate_norm), log)
print_log("Distance Pruning Rate: {}".format(args.rate_dist), log)
print_log("Layer Begin: {}".format(args.layer_begin), log)
print_log("Layer End: {}".format(args.layer_end), log)
print_log("Layer Inter: {}".format(args.layer_inter), log)
print_log("Epoch prune: {}".format(args.epoch_prune), log)
print_log("use pretrain: {}".format(args.use_pretrain), log)
print_log("Pretrain path: {}".format(args.pretrain_path), log)
print_log("Dist type: {}".format(args.dist_type), log)
print_log("Pre cfg: {}".format(args.use_precfg), log)
if args.dataset == 'cifar10':
train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(
args.data_path,
train=True,
download=True,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR10(args.data_path,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])),
batch_size=args.test_batch_size,
shuffle=False,
**kwargs)
else:
train_loader = torch.utils.data.DataLoader(datasets.CIFAR100(
args.data_path,
train=True,
download=True,
transform=transforms.Compose([
transforms.Pad(4),
transforms.RandomCrop(32),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])),
batch_size=args.batch_size,
shuffle=True,
**kwargs)
test_loader = torch.utils.data.DataLoader(
datasets.CIFAR100(args.data_path,
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(
(0.4914, 0.4822, 0.4465),
(0.2023, 0.1994, 0.2010))
])),
batch_size=args.test_batch_size,
shuffle=True,
**kwargs)
print_log("=> creating model '{}'".format(args.arch), log)
model = models.__dict__[args.arch](dataset=args.dataset, depth=args.depth)
print_log("=> network :\n {}".format(model), log)
if args.cuda:
model.cuda()
if args.use_pretrain:
if os.path.isfile(args.pretrain_path):
print_log(
"=> loading pretrain model '{}'".format(args.pretrain_path),
log)
else:
dir = '/home/yahe/compress/filter_similarity/logs/main_2'
args.pretrain_path = dir + '/checkpoint.pth.tar'
print_log("Pretrain path: {}".format(args.pretrain_path), log)
pretrain = torch.load(args.pretrain_path)
if args.use_state_dict:
model.load_state_dict(pretrain['state_dict'])
else:
model = pretrain['state_dict']
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model = models.vgg(dataset='cifar10',
depth=16,
cfg=checkpoint['cfg'])
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {}) Prec1: {:f}".format(
args.resume, checkpoint['epoch'], best_prec1))
if args.cuda:
model = model.cuda()
else:
print("=> no checkpoint found at '{}'".format(args.resume))
if args.evaluate:
time1 = time.time()
test(test_loader, model, log)
time2 = time.time()
print('function took %0.3f ms' % ((time2 - time1) * 1000.0))
return
m = Mask(model)
m.init_length()
print("-" * 10 + "one epoch begin" + "-" * 10)
print("remaining ratio of pruning : Norm is %f" % args.rate_norm)
print("reducing ratio of pruning : Distance is %f" % args.rate_dist)
print("total remaining ratio is %f" % (args.rate_norm - args.rate_dist))
val_acc_1 = test(test_loader, model, log)
print(" accu before is: %.3f %%" % val_acc_1)
m.model = model
m.init_mask(args.rate_norm, args.rate_dist, args.dist_type)
# m.if_zero()
m.do_mask()
m.do_similar_mask()
model = m.model
# m.if_zero()
if args.cuda:
model = model.cuda()
val_acc_2 = test(test_loader, model, log)
print(" accu after is: %s %%" % val_acc_2)
best_prec1 = 0.
for epoch in range(args.start_epoch, args.epochs):
if epoch in [args.epochs * 0.5, args.epochs * 0.75]:
for param_group in optimizer.param_groups:
param_group['lr'] *= 0.1
train(train_loader, model, optimizer, epoch, log)
prec1 = test(test_loader, model, log)
if epoch % args.epoch_prune == 0 or epoch == args.epochs - 1:
m.model = model
m.if_zero()
m.init_mask(args.rate_norm, args.rate_dist, args.dist_type)
m.do_mask()
m.do_similar_mask()
# small_filter_index.append(m.filter_small_index)
# large_filter_index.append(m.filter_large_index)
# save_obj(small_filter_index, 'small_filter_index_2')
# save_obj(large_filter_index, 'large_filter_index_2')
m.if_zero()
model = m.model
if args.cuda:
model = model.cuda()
val_acc_2 = test(test_loader, model, log)
is_best = val_acc_2 > best_prec1
best_prec1 = max(prec1, best_prec1)
save_checkpoint(
{
'epoch': epoch + 1,
'state_dict': model.state_dict(),
'best_prec1': best_prec1,
'optimizer': optimizer.state_dict(),
'cfg': model.cfg
},
is_best,
filepath=args.save_path)
class LossHistory(Callback):
def on_train_begin(self, logs={}):
self.losses = []
def on_batch_end(self, batch, logs={}):
self.losses.append(logs.get('loss'))
history = LossHistory()
batch_size = 300
num_epochs = 50
print("\n")
model = vgg()
model.fit(X_train, y_train,
batch_size=batch_size,
epochs=num_epochs,
verbose=1,
validation_data=(X_test, y_test),
callbacks=[tb, history])
print('Saving the model now')
model.save('vgg.h5')
print('Done saving the model')
del model;
model = load_model('vgg.h5')
model.summary()
score = model.evaluate(X_test, y_test, verbose=1)
print("-" * 40)
print("VGG Model (%d epochs):" % num_epochs)
import tflearn
from models import vgg
from data import test_x, test_y, categories
vgg_model = tflearn.DNN(vgg(len(categories)))
vgg_model.load('../trained_models/vgg/vgg_model')
accuracy = vgg_model.evaluate(test_x, test_y)
print('VGG accuracy: ', accuracy)
type=str,
metavar='PATH',
help='path to the model (default: none)')
parser.add_argument('--save',
default='./logs',
type=str,
metavar='PATH',
help='path to save pruned model (default: none)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
args.cuda = False
if not os.path.exists(args.save):
os.makedirs(args.save)
model = vgg(dataset=args.dataset, depth=args.depth)
if args.cuda:
model.cuda()
if args.model:
if os.path.isfile(args.model):
print("=> loading checkpoint '{}'".format(args.model))
checkpoint = torch.load(args.model)
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {}) Prec1: {:f}".format(
args.model, checkpoint['epoch'], best_prec1))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
help='path to the model (default: none)')
parser.add_argument('--save',
default='.',
type=str,
metavar='PATH',
help='path to save pruned model (default: none)')
args = parser.parse_args()
args.cuda = not args.no_cuda and torch.cuda.is_available()
if not os.path.exists(args.save):
os.makedirs(args.save)
#model = vgg(dataset=args.dataset, depth=args.depth)
model = vgg(num_channel=args.num_channel,
dataset=args.dataset,
depth=args.depth)
#ckpt_file = join(args.model)
if args.cuda:
model.cuda()
if args.model:
print(args.model)
if os.path.isfile(args.model):
print("=> loading checkpoint '{}'".format(args.model))
checkpoint = torch.load(
args.model) #, map_location = torch.device('cpu'))
args.start_epoch = checkpoint['epoch']
best_prec1 = checkpoint['best_prec1']
try:
m, nn.MaxPool2D) or isinstance(m, nn.AvgPool2D):
m.register_forward_hook(pooling_hook)
if isinstance(m, nn.HybridSequential):
get(m)
get(model)
input = nd.random.uniform(-1,
1,
shape=(1, 3, input_res, input_res),
ctx=mx.cpu(0))
model.initialize(mxnet.init.Xavier(), ctx=mx.cpu(0), force_reinit=True)
out = model(input)
total_flops = (sum(list_conv_flops) + sum(list_dense_flops) +
sum(list_bn_flops) + sum(list_relu) + sum(list_pooling))
total_params = (sum(list_conv_params) + sum(list_dense_params) +
sum(list_bn_params))
print('Number of params: %.2fM' % (total_params / 1e6))
print('\nNumber of FLOPs: %.5fG' % (total_flops / 3 / 1e9))
return total_flops / 3, total_params
if __name__ == '__main__':
import models
net = models.vgg()
#print(net)
flops, _ = print_model_param_flops(net, input_res=32, multiply_adds=True)
#print(flops)
encoder_c, encoder_t = torch.load(enc_c_name), torch.load(enc_c_name)
optim_enc_t = optim.Adam(encoder_t.parameters(), lr=args.lr, betas=(0.5, 0.9))
# Discriminator
if args.gan_ratio != 0:
t_discriminator_patch = models.Discriminator_patch(channels=3).cuda()
optim_disc_patch = optim.Adam(filter(lambda p: p.requires_grad, t_discriminator_patch.parameters()), lr=args.lr, betas=(0.0,0.9))
# Relation
if args.relation_ratio != 0:
R = models.Relation().cuda()
R.set(encoder_c)
optim_R = optim.Adam(R.model_bn1_B.parameters(), lr=args.lr, betas=(0.5, 0.9))
# Feature Extractor
F = models.vgg().cuda().eval()
### Lab shuffle
lab = []
def preprocessing(data):
new_data = data.clone()
global iters, lab
# Record min/max a/b
if iters == 0:
max_a, min_a, max_b, min_b = -10000, 10000, -10000, 10000
for b in range(new_data.size(0)):
data2 = new_data[b].cpu().data.numpy().transpose(1,2,0).astype(float) #*255
data2 = color.rgb2lab(data2)
max_a, min_a = max(max_a, np.max(data2[:,:,1])), min(min_a, np.min(data2[:,:,1]))
max_b, min_b = max(max_b, np.max(data2[:,:,2])), min(min_b, np.min(data2[:,:,2]))
lab = [[min_a, max_a], [min_b, max_b]]
import tflearn, os
from models import vgg
from data import train_x, train_y, test_x, test_y, categories
# Create directories if necessary
save_path = os.path.dirname(
os.path.realpath(__file__)) + '/../trained_models/simple'
if not os.path.exists(save_path):
os.makedirs(save_path)
vgg_model = tflearn.DNN(
vgg(len(categories)),
tensorboard_verbose=3,
tensorboard_dir='../logs',
)
vgg_model.fit({'input': train_x}, {'targets': train_y},
n_epoch=100,
validation_set=({
'input': test_x
}, {
'targets': test_y
}),
batch_size=64,
snapshot_step=200,
show_metric=True,
run_id='vgg')
vgg_model.save('../trained_models/vgg/vgg_model')
