未验证 提交 4fbfc011 编写于 作者: W whs 提交者: GitHub

Restore quantization and distillation stategy before loading persistables. (#16959)

test=develop
上级 5a2d6d6b
......@@ -363,6 +363,9 @@ class Compressor(object):
strategies = pickle.load(
strategy_file, encoding='bytes')
for strategy in strategies:
strategy.restore_from_checkpoint(context)
if os.path.exists(model_path):
exe = SlimGraphExecutor(context.place)
with scope_guard(context.scope):
......
......@@ -46,3 +46,6 @@ class Strategy(object):
def on_compression_end(self, context):
pass
def restore_from_checkpoint(self, context):
pass
......@@ -38,7 +38,7 @@ class DistillationStrategy(Strategy):
super(DistillationStrategy, self).__init__(start_epoch, end_epoch)
self.distillers = distillers
def on_compression_begin(self, context):
def restore_from_checkpoint(self, context):
# load from checkpoint
if context.epoch_id > 0:
if context.epoch_id > self.start_epoch and context.epoch_id < self.end_epoch:
......
......@@ -88,7 +88,7 @@ class QuantizationStrategy(Strategy):
self.save_out_nodes = save_out_nodes
self.save_in_nodes = save_in_nodes
def on_compression_begin(self, context):
def restore_from_checkpoint(self, context):
"""
Restore graph when the compressoin task is inited from checkpoint.
"""
......@@ -143,10 +143,9 @@ class QuantizationStrategy(Strategy):
train_ir_graph.graph).with_data_parallel(
loss_name=context.optimize_graph.out_nodes['loss'],
build_strategy=build_strategy)
# for evaluation. And program compiled from ir graph must be with data parallel.
context.eval_graph.compiled_graph = CompiledProgram(
test_ir_graph.graph).with_data_parallel(
build_strategy=build_strategy)
context.eval_graph.program = test_ir_graph.to_program()
# for saving inference model after training
context.put('quantization_test_ir_graph_backup', test_ir_graph)
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册