def save_prune_checkpoint(exe, program, ckpt_name):
"""
Save checkpoint for evaluation or resume training
"""
ckpt_dir = os.path.join(cfg.TRAIN.MODEL_SAVE_DIR, str(ckpt_name))
print("Save model checkpoint to {}".format(ckpt_dir))
if not os.path.isdir(ckpt_dir):
os.makedirs(ckpt_dir)
save_model(exe, program, ckpt_dir)
return ckpt_dir
def test_prune(self):
train_program = fluid.Program()
startup_program = fluid.Program()
with fluid.program_guard(train_program, startup_program):
input = fluid.data(name="image", shape=[None, 3, 16, 16])
conv1 = conv_bn_layer(input, 8, 3, "conv1")
conv2 = conv_bn_layer(conv1, 8, 3, "conv2")
sum1 = conv1 + conv2
conv3 = conv_bn_layer(sum1, 8, 3, "conv3")
conv4 = conv_bn_layer(conv3, 8, 3, "conv4")
sum2 = conv4 + sum1
conv5 = conv_bn_layer(sum2, 8, 3, "conv5")
conv6 = conv_bn_layer(conv5, 8, 3, "conv6")
feature = fluid.layers.reshape(conv6, [-1, 128, 16])
predict = fluid.layers.fc(input=feature, size=10, act='softmax')
label = fluid.data(name='label', shape=[None, 1], dtype='int64')
print(label.shape)
print(predict.shape)
cost = fluid.layers.cross_entropy(input=predict, label=label)
avg_cost = fluid.layers.mean(cost)
adam_optimizer = fluid.optimizer.AdamOptimizer(0.01)
adam_optimizer.minimize(avg_cost)
place = fluid.CPUPlace()
exe = fluid.Executor(place)
scope = fluid.global_scope()
exe.run(startup_program, scope=scope)
criterion = 'bn_scale'
pruner = Pruner(criterion)
main_program, _, _ = pruner.prune(train_program,
scope,
params=["conv4_weights"],
ratios=[0.5],
place=place,
lazy=False,
only_graph=False,
param_backup=None,
param_shape_backup=None)
x = numpy.random.random(size=(10, 3, 16, 16)).astype('float32')
label = numpy.random.random(size=(10, 1)).astype('int64')
loss_data, = exe.run(train_program,
feed={
"image": x,
"label": label
},
fetch_list=[cost.name])
save_model(exe, main_program, 'model_file')
pruned_program = fluid.Program()
pruned_startup_program = fluid.Program()
with fluid.program_guard(pruned_program, pruned_startup_program):
input = fluid.data(name="image", shape=[None, 3, 16, 16])
conv1 = conv_bn_layer(input, 8, 3, "conv1")
conv2 = conv_bn_layer(conv1, 8, 3, "conv2")
sum1 = conv1 + conv2
conv3 = conv_bn_layer(sum1, 8, 3, "conv3")
conv4 = conv_bn_layer(conv3, 8, 3, "conv4")
sum2 = conv4 + sum1
conv5 = conv_bn_layer(sum2, 8, 3, "conv5")
conv6 = conv_bn_layer(conv5, 8, 3, "conv6")
pruned_test_program = pruned_program.clone(for_test=True)
exe.run(pruned_startup_program)
load_model(exe, pruned_program, 'model_file')
load_model(exe, pruned_test_program, 'model_file')
shapes = {
"conv1_weights": (4, 3, 3, 3),
"conv2_weights": (4, 4, 3, 3),
"conv3_weights": (8, 4, 3, 3),
"conv4_weights": (4, 8, 3, 3),
"conv5_weights": (8, 4, 3, 3),
"conv6_weights": (8, 8, 3, 3)
}
for param in pruned_program.global_block().all_parameters():
if "weights" in param.name:
print("param: {}; param shape: {}".format(
param.name, param.shape))
self.assertTrue(param.shape == shapes[param.name])
for param in pruned_test_program.global_block().all_parameters():
if "weights" in param.name:
print("param: {}; param shape: {}".format(
param.name, param.shape))
self.assertTrue(param.shape == shapes[param.name])
def compress(args):
train_reader = None
test_reader = None
if args.data == "mnist":
import paddle.dataset.mnist as reader
train_reader = reader.train()
val_reader = reader.test()
class_dim = 10
image_shape = "1,28,28"
elif args.data == "imagenet":
import imagenet_reader as reader
train_reader = reader.train()
val_reader = reader.val()
class_dim = 1000
image_shape = "3,224,224"
else:
raise ValueError("{} is not supported.".format(args.data))
image_shape = [int(m) for m in image_shape.split(",")]
assert args.model in model_list, "{} is not in lists: {}".format(
args.model, model_list)
image = fluid.layers.data(name='image', shape=image_shape, dtype='float32')
label = fluid.layers.data(name='label', shape=[1], dtype='int64')
# model definition
model = models.__dict__[args.model]()
out = model.net(input=image, class_dim=class_dim)
cost = fluid.layers.cross_entropy(input=out, label=label)
avg_cost = fluid.layers.mean(x=cost)
acc_top1 = fluid.layers.accuracy(input=out, label=label, k=1)
acc_top5 = fluid.layers.accuracy(input=out, label=label, k=5)
val_program = fluid.default_main_program().clone(for_test=True)
opt = create_optimizer(args)
opt.minimize(avg_cost)
place = fluid.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
exe = fluid.Executor(place)
exe.run(fluid.default_startup_program())
if args.pretrained_model:
def if_exist(var):
return os.path.exists(os.path.join(args.pretrained_model,
var.name))
_logger.info("Load pretrained model from {}".format(
args.pretrained_model))
fluid.io.load_vars(exe, args.pretrained_model, predicate=if_exist)
val_reader = paddle.fluid.io.batch(val_reader, batch_size=args.batch_size)
train_reader = paddle.fluid.io.batch(train_reader,
batch_size=args.batch_size,
drop_last=True)
train_feeder = feeder = fluid.DataFeeder([image, label], place)
val_feeder = feeder = fluid.DataFeeder([image, label],
place,
program=val_program)
def test(epoch, program):
batch_id = 0
acc_top1_ns = []
acc_top5_ns = []
for data in val_reader():
start_time = time.time()
acc_top1_n, acc_top5_n = exe.run(
program,
feed=train_feeder.feed(data),
fetch_list=[acc_top1.name, acc_top5.name])
end_time = time.time()
if batch_id % args.log_period == 0:
_logger.info(
"Eval epoch[{}] batch[{}] - acc_top1: {}; acc_top5: {}; time: {}"
.format(epoch, batch_id, np.mean(acc_top1_n),
np.mean(acc_top5_n), end_time - start_time))
acc_top1_ns.append(np.mean(acc_top1_n))
acc_top5_ns.append(np.mean(acc_top5_n))
batch_id += 1
_logger.info(
"Final eval epoch[{}] - acc_top1: {}; acc_top5: {}".format(
epoch, np.mean(np.array(acc_top1_ns)),
np.mean(np.array(acc_top5_ns))))
def train(epoch, program):
build_strategy = fluid.BuildStrategy()
exec_strategy = fluid.ExecutionStrategy()
train_program = fluid.compiler.CompiledProgram(
program).with_data_parallel(loss_name=avg_cost.name,
build_strategy=build_strategy,
exec_strategy=exec_strategy)
batch_id = 0
for data in train_reader():
start_time = time.time()
loss_n, acc_top1_n, acc_top5_n = exe.run(
train_program,
feed=train_feeder.feed(data),
fetch_list=[avg_cost.name, acc_top1.name, acc_top5.name])
end_time = time.time()
loss_n = np.mean(loss_n)
acc_top1_n = np.mean(acc_top1_n)
acc_top5_n = np.mean(acc_top5_n)
if batch_id % args.log_period == 0:
_logger.info(
"epoch[{}]-batch[{}] - loss: {}; acc_top1: {}; acc_top5: {}; time: {}"
.format(epoch, batch_id, loss_n, acc_top1_n, acc_top5_n,
end_time - start_time))
batch_id += 1
test(0, val_program)
params = get_pruned_params(args, fluid.default_main_program())
_logger.info("FLOPs before pruning: {}".format(
flops(fluid.default_main_program())))
pruner = Pruner(args.criterion)
pruned_val_program, _, _ = pruner.prune(val_program,
fluid.global_scope(),
params=params,
ratios=[args.pruned_ratio] *
len(params),
place=place,
only_graph=True)
pruned_program, _, _ = pruner.prune(fluid.default_main_program(),
fluid.global_scope(),
params=params,
ratios=[args.pruned_ratio] *
len(params),
place=place)
_logger.info("FLOPs after pruning: {}".format(flops(pruned_program)))
for i in range(args.num_epochs):
train(i, pruned_program)
if i % args.test_period == 0:
test(i, pruned_val_program)
save_model(exe, pruned_val_program,
os.path.join(args.model_path, str(i)))
if args.save_inference:
infer_model_path = os.path.join(args.model_path, "infer_models",
str(i))
fluid.io.save_inference_model(infer_model_path, ["image"], [out],
exe, pruned_val_program)
_logger.info(
"Saved inference model into [{}]".format(infer_model_path))
def compress(args):
train_reader = None
test_reader = None
if args.data == "mnist":
transform = T.Compose([T.Transpose(), T.Normalize([127.5], [127.5])])
train_dataset = paddle.vision.datasets.MNIST(mode='train',
backend="cv2",
transform=transform)
val_dataset = paddle.vision.datasets.MNIST(mode='test',
backend="cv2",
transform=transform)
class_dim = 10
image_shape = "1,28,28"
elif args.data == "imagenet":
import imagenet_reader as reader
train_dataset = reader.ImageNetDataset(mode='train')
val_dataset = reader.ImageNetDataset(mode='val')
class_dim = 1000
image_shape = "3,224,224"
else:
raise ValueError("{} is not supported.".format(args.data))
image_shape = [int(m) for m in image_shape.split(",")]
assert args.model in model_list, "{} is not in lists: {}".format(
args.model, model_list)
places = paddle.static.cuda_places(
) if args.use_gpu else paddle.static.cpu_places()
place = places[0]
exe = paddle.static.Executor(place)
image = paddle.static.data(name='image',
shape=[None] + image_shape,
dtype='float32')
label = paddle.static.data(name='label', shape=[None, 1], dtype='int64')
batch_size_per_card = int(args.batch_size / len(places))
train_loader = paddle.io.DataLoader(train_dataset,
places=places,
feed_list=[image, label],
drop_last=True,
batch_size=batch_size_per_card,
shuffle=True,
return_list=False,
use_shared_memory=True,
num_workers=16)
valid_loader = paddle.io.DataLoader(val_dataset,
places=place,
feed_list=[image, label],
drop_last=False,
return_list=False,
use_shared_memory=True,
batch_size=batch_size_per_card,
shuffle=False)
step_per_epoch = int(np.ceil(len(train_dataset) * 1. / args.batch_size))
# model definition
model = models.__dict__[args.model]()
out = model.net(input=image, class_dim=class_dim)
cost = paddle.nn.functional.loss.cross_entropy(input=out, label=label)
avg_cost = paddle.mean(x=cost)
acc_top1 = paddle.metric.accuracy(input=out, label=label, k=1)
acc_top5 = paddle.metric.accuracy(input=out, label=label, k=5)
val_program = paddle.static.default_main_program().clone(for_test=True)
opt, learning_rate = create_optimizer(args, step_per_epoch)
opt.minimize(avg_cost)
exe.run(paddle.static.default_startup_program())
if args.pretrained_model:
def if_exist(var):
return os.path.exists(os.path.join(args.pretrained_model,
var.name))
_logger.info("Load pretrained model from {}".format(
args.pretrained_model))
paddle.static.load(paddle.static.default_main_program(),
args.pretrained_model, exe)
def test(epoch, program):
acc_top1_ns = []
acc_top5_ns = []
for batch_id, data in enumerate(valid_loader):
start_time = time.time()
acc_top1_n, acc_top5_n = exe.run(
program, feed=data, fetch_list=[acc_top1.name, acc_top5.name])
end_time = time.time()
if batch_id % args.log_period == 0:
_logger.info(
"Eval epoch[{}] batch[{}] - acc_top1: {}; acc_top5: {}; time: {}"
.format(epoch, batch_id, np.mean(acc_top1_n),
np.mean(acc_top5_n), end_time - start_time))
acc_top1_ns.append(np.mean(acc_top1_n))
acc_top5_ns.append(np.mean(acc_top5_n))
_logger.info(
"Final eval epoch[{}] - acc_top1: {}; acc_top5: {}".format(
epoch, np.mean(np.array(acc_top1_ns)),
np.mean(np.array(acc_top5_ns))))
def train(epoch, program):
for batch_id, data in enumerate(train_loader):
start_time = time.time()
loss_n, acc_top1_n, acc_top5_n = exe.run(
train_program,
feed=data,
fetch_list=[avg_cost.name, acc_top1.name, acc_top5.name])
end_time = time.time()
loss_n = np.mean(loss_n)
acc_top1_n = np.mean(acc_top1_n)
acc_top5_n = np.mean(acc_top5_n)
if batch_id % args.log_period == 0:
_logger.info(
"epoch[{}]-batch[{}] lr: {:.6f} - loss: {}; acc_top1: {}; acc_top5: {}; time: {}"
.format(epoch, batch_id, learning_rate.get_lr(), loss_n,
acc_top1_n, acc_top5_n, end_time - start_time))
learning_rate.step()
batch_id += 1
test(0, val_program)
params = get_pruned_params(args, paddle.static.default_main_program())
_logger.info("FLOPs before pruning: {}".format(
flops(paddle.static.default_main_program())))
pruner = Pruner(args.criterion)
pruned_val_program, _, _ = pruner.prune(val_program,
paddle.static.global_scope(),
params=params,
ratios=[args.pruned_ratio] *
len(params),
place=place,
only_graph=True)
pruned_program, _, _ = pruner.prune(paddle.static.default_main_program(),
paddle.static.global_scope(),
params=params,
ratios=[args.pruned_ratio] *
len(params),
place=place)
_logger.info("FLOPs after pruning: {}".format(flops(pruned_program)))
build_strategy = paddle.static.BuildStrategy()
exec_strategy = paddle.static.ExecutionStrategy()
train_program = paddle.static.CompiledProgram(
pruned_program).with_data_parallel(loss_name=avg_cost.name,
build_strategy=build_strategy,
exec_strategy=exec_strategy)
for i in range(args.num_epochs):
train(i, train_program)
if (i + 1) % args.test_period == 0:
test(i, pruned_val_program)
save_model(exe, pruned_val_program,
os.path.join(args.model_path, str(i)))
if args.save_inference:
infer_model_path = os.path.join(args.model_path, "infer_models",
str(i))
paddle.static.save_inference_model(infer_model_path, [image],
[out],
exe,
program=pruned_val_program)
_logger.info(
"Saved inference model into [{}]".format(infer_model_path))
def test_prune(self):
train_program = fluid.Program()
startup_program = fluid.Program()
with fluid.program_guard(train_program, startup_program):
input = fluid.data(name="image", shape=[None, 3, 16, 16])
conv1 = conv_bn_layer(input, 8, 3, "conv1")
conv2 = conv_bn_layer(conv1, 8, 3, "conv2")
sum1 = conv1 + conv2
conv3 = conv_bn_layer(sum1, 8, 3, "conv3")
conv4 = conv_bn_layer(conv3, 8, 3, "conv4")
sum2 = conv4 + sum1
conv5 = conv_bn_layer(sum2, 8, 3, "conv5")
conv6 = conv_bn_layer(conv5, 8, 3, "conv6")
place = fluid.CPUPlace()
exe = fluid.Executor(place)
scope = fluid.global_scope()
exe.run(startup_program, scope=scope)
criterion = 'bn_scale'
pruner = Pruner(criterion)
main_program, _, _ = pruner.prune(train_program,
scope,
params=["conv4_weights"],
ratios=[0.5],
place=place,
lazy=False,
only_graph=False,
param_backup=None,
param_shape_backup=None)
x = numpy.random.random(size=(10, 3, 16, 16)).astype('float32')
loss_data, = exe.run(train_program,
feed={"image": x},
fetch_list=[conv6.name])
save_model(exe, main_program, 'model_file')
pruned_program = fluid.Program()
pruned_startup_program = fluid.Program()
with fluid.program_guard(pruned_program, pruned_startup_program):
input = fluid.data(name="image", shape=[None, 3, 16, 16])
conv1 = conv_bn_layer(input, 8, 3, "conv1")
conv2 = conv_bn_layer(conv1, 8, 3, "conv2")
sum1 = conv1 + conv2
conv3 = conv_bn_layer(sum1, 8, 3, "conv3")
conv4 = conv_bn_layer(conv3, 8, 3, "conv4")
sum2 = conv4 + sum1
conv5 = conv_bn_layer(sum2, 8, 3, "conv5")
conv6 = conv_bn_layer(conv5, 8, 3, "conv6")
exe.run(pruned_startup_program)
load_model(exe, pruned_program, 'model_file')
shapes = {
"conv1_weights": (4, 3, 3, 3),
"conv2_weights": (4, 4, 3, 3),
"conv3_weights": (8, 4, 3, 3),
"conv4_weights": (4, 8, 3, 3),
"conv5_weights": (8, 4, 3, 3),
"conv6_weights": (8, 8, 3, 3)
}
for param in pruned_program.global_block().all_parameters():
if "weights" in param.name:
print("param: {}; param shape: {}".format(
param.name, param.shape))
self.assertTrue(param.shape == shapes[param.name])