class TestAdaptorONNXRT(unittest.TestCase):
mb_v2_export_path = "mb_v2.onnx"
mb_v2_model = torchvision.models.mobilenet_v2()
rn50_export_path = "rn50.onnx"
rn50_model = torchvision.models.resnet50()
datasets = DATASETS('onnxrt_qlinearops')
cv_dataset = datasets['dummy'](shape=(100, 3, 224, 224), low=0., high=1., label=True)
cv_dataloader = DATALOADERS['onnxrt_qlinearops'](cv_dataset)
@classmethod
def setUpClass(self):
build_static_yaml()
build_dynamic_yaml()
build_non_MSE_yaml()
export_onnx_model(self.mb_v2_model, self.mb_v2_export_path)
self.mb_v2_model = onnx.load(self.mb_v2_export_path)
export_onnx_model(self.rn50_model, self.rn50_export_path)
self.rn50_model = onnx.load(self.rn50_export_path)
@classmethod
def tearDownClass(self):
os.remove("static_yaml.yaml")
os.remove("dynamic_yaml.yaml")
os.remove("non_MSE_yaml.yaml")
os.remove(self.mb_v2_export_path)
os.remove(self.rn50_export_path)
shutil.rmtree("./saved", ignore_errors=True)
shutil.rmtree("runs", ignore_errors=True)
def test_adaptor(self):
framework_specific_info = {"device": "cpu",
"approach": "post_training_static_quant",
"random_seed": 1234,
"q_dataloader": None,
"backend": "qlinearops",
"workspace_path": './lpot_workspace/{}/{}/'.format(
'onnxrt',
'imagenet')}
framework = "onnxrt_qlinearops"
adaptor = FRAMEWORKS[framework](framework_specific_info)
adaptor.inspect_tensor(self.rn50_model, self.cv_dataloader, ["Conv"])
def test_quantizate(self):
from lpot.experimental import Quantization, common
for fake_yaml in ["static_yaml.yaml", "dynamic_yaml.yaml"]:
quantizer = Quantization(fake_yaml)
quantizer.calib_dataloader = self.cv_dataloader
quantizer.eval_dataloader = self.cv_dataloader
quantizer.model = common.Model(self.rn50_model)
q_model = quantizer()
eval_func(q_model)
for fake_yaml in ["non_MSE_yaml.yaml"]:
quantizer = Quantization(fake_yaml)
quantizer.calib_dataloader = self.cv_dataloader
quantizer.eval_dataloader = self.cv_dataloader
quantizer.model = common.Model(self.mb_v2_model)
q_model = quantizer()
eval_func(q_model)
def test_tensorflow_dummy(self):
datasets = DATASETS('tensorflow')
dataset = datasets['dummy'](shape=(4, 256, 256, 3))
data_loader = DATALOADERS['tensorflow'](dataset)
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
with self.assertRaises(AssertionError):
dataset = datasets['dummy'](shape=[(4, 256, 256, 3), (256, 256,
3)])
with self.assertRaises(AssertionError):
dataset = datasets['dummy'](shape=(4, 256, 256, 3), low=[1., 0.])
with self.assertRaises(AssertionError):
dataset = datasets['dummy'](shape=(4, 256, 256, 3),
high=[128., 127.])
with self.assertRaises(AssertionError):
dataset = datasets['dummy'](shape=(4, 256, 256, 3),
dtype=['float32', 'int8'])
def test_footprint(self):
from lpot.experimental import Benchmark, common
from lpot.data import DATASETS
dataset = DATASETS('tensorflow')['dummy']((100, 256, 256, 1),
label=True)
benchmarker = Benchmark('fake_yaml_footprint.yaml')
benchmarker.b_dataloader = common.DataLoader(dataset)
benchmarker.model = self.constant_graph_1
benchmarker()
def test_style_transfer_dataset(self):
jpg_url = "https://upload.wikimedia.org/wikipedia/commons/thumb/c/c1/Doll_face_silver_Persian.jpg/1024px-Doll_face_silver_Persian.jpg"
os.system("wget {} -O test.jpg".format(jpg_url))
datasets = DATASETS('tensorflow')
dataset = datasets['style_transfer'](content_folder='./',
style_folder='./')
length = len(dataset)
image, label = dataset[0]
self.assertEqual(image[0].shape, (256, 256, 3))
self.assertEqual(image[1].shape, (256, 256, 3))
os.remove('test.jpg')
def test_onnx_integer_dummy(self):
datasets = DATASETS('onnxrt_integerops')
dataset = datasets['dummy'](shape=(4, 256, 256, 3))
data_loader = DATALOADERS['onnxrt_integerops'](dataset)
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
def test_mxnet_dummy(self):
datasets = DATASETS('mxnet')
dataset = datasets['dummy'](shape=(4, 256, 256, 3))
data_loader = DataLoader('mxnet', dataset)
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
def test_style_transfer_dataset(self):
random_array = np.random.random_sample([100, 100, 3]) * 255
random_array = random_array.astype(np.uint8)
im = Image.fromarray(random_array)
im.save('test.jpg')
datasets = DATASETS('tensorflow')
dataset = datasets['style_transfer'](content_folder='./',
style_folder='./')
length = len(dataset)
image, label = dataset[0]
self.assertEqual(image[0].shape, (256, 256, 3))
self.assertEqual(image[1].shape, (256, 256, 3))
os.remove('test.jpg')
def test_pytorch_dummy(self):
datasets = DATASETS('pytorch')
dataset = datasets['dummy'](shape=[(4, 256, 256, 3), (4, 1)], \
high=[10., 10.], low=[0., 0.])
data_loader = DataLoader('pytorch', dataset)
iterator = iter(data_loader)
data, label = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data, label = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
def test_pytorch_dummy(self):
datasets = DATASETS('pytorch')
transform = TRANSFORMS('pytorch', 'preprocess')['Resize'](**{
'size': 100
})
dataset = datasets['dummy'](shape=[(4, 256, 256, 3), (4, 1)], \
high=[10., 10.], low=[0., 0.], transform=transform)
data_loader = DATALOADERS['pytorch'](dataset)
iterator = iter(data_loader)
data, label = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data, label = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
def test_onnxrt_qlinear_dummy(self):
datasets = DATASETS('onnxrt_qlinearops')
transform = TRANSFORMS('onnxrt_qlinearops',
'preprocess')['Resize'](**{
'size': 100
})
dataset = datasets['dummy'](shape=(4, 256, 256, 3),
transform=transform)
data_loader = DATALOADERS['onnxrt_qlinearops'](dataset)
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
def test_performance(self):
from lpot.data import DATASETS
dataset = DATASETS('tensorflow')['dummy']((100, 256, 256, 1),
label=True)
from lpot.experimental import Quantization, common
from lpot.utils.utility import get_size
quantizer = Quantization('fake_yaml.yaml')
quantizer.calib_dataloader = common.DataLoader(dataset)
quantizer.eval_dataloader = common.DataLoader(dataset)
quantizer.model = self.constant_graph
q_model = quantizer()
from lpot.experimental import Benchmark, common
benchmarker = Benchmark('fake_yaml.yaml')
benchmarker.b_dataloader = common.DataLoader(dataset)
benchmarker.model = self.constant_graph_1
benchmarker()
def create_nlp_session():
a_value = np.random.randn(100, 4).astype(np.float32)
A_init = helper.make_tensor('A', TensorProto.FLOAT, [100, 4],
a_value.reshape(400).tolist())
b_value = np.random.randint(2, size=(10)).astype(np.int32)
B_init = helper.make_tensor('B', TensorProto.INT32, [10],
b_value.reshape(10).tolist())
A = helper.make_tensor_value_info('A', TensorProto.FLOAT, [100, 4])
B = helper.make_tensor_value_info('B', TensorProto.INT32, [10])
C = helper.make_tensor_value_info('C', TensorProto.FLOAT, [10, 4])
node = onnx.helper.make_node('Gather', ['A', 'B'], ['C'], name='Gather')
graph = helper.make_graph([node], 'test_graph_1', [A, B], [C],
[A_init, B_init])
model = helper.make_model(graph)
datasets = DATASETS('onnxrt_qlinearops')
dataset = datasets['dummy'](shape=(100, 4), label=True)
dataloader = DATALOADERS['onnxrt_qlinearops'](dataset)
return model, dataloader
def test_pruning_external(self):
from lpot.experimental import common
from lpot import Pruning
prune = Pruning('fake.yaml')
datasets = DATASETS('pytorch')
dummy_dataset = datasets['dummy'](shape=(100, 3, 224, 224),
low=0.,
high=1.,
label=True)
dummy_dataloader = PyTorchDataLoader(dummy_dataset)
def training_func_for_lpot(model):
epochs = 16
iters = 30
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
for nepoch in range(epochs):
model.train()
cnt = 0
prune.on_epoch_begin(nepoch)
for image, target in dummy_dataloader:
prune.on_batch_begin(cnt)
print('.', end='')
cnt += 1
output = model(image)
loss = criterion(output, target)
optimizer.zero_grad()
loss.backward()
optimizer.step()
prune.on_batch_end()
if cnt >= iters:
break
prune.on_epoch_end()
prune.model = common.Model(self.model)
prune.pruning_func = training_func_for_lpot
prune.eval_dataloader = dummy_dataloader
prune.train_dataloader = dummy_dataloader
_ = prune(common.Model(self.model), \
train_dataloader=dummy_dataloader, \
pruning_func=training_func_for_lpot, \
eval_dataloader=dummy_dataloader)
def create_cv_session():
A = helper.make_tensor_value_info('A', TensorProto.FLOAT, [1, 1, 5, 5])
B = helper.make_tensor_value_info('B', TensorProto.FLOAT, [1, 1, 3, 3])
b_value = np.random.randn(1, 1, 3, 3).astype(np.float32)
B_init = helper.make_tensor('B', TensorProto.FLOAT, [1, 1, 3, 3],
b_value.reshape(9).tolist())
D = helper.make_tensor_value_info('D', TensorProto.FLOAT, [1, 1, 5, 5])
conv_node = onnx.helper.make_node('Conv', ['A', 'B'], ['C'],
name='Conv',
kernel_shape=[3, 3],
pads=[1, 1, 1, 1])
relu_node = onnx.helper.make_node('Relu', ['C'], ['D'], name='Relu')
graph = helper.make_graph([conv_node, relu_node], 'test_graph_1', [A, B],
[D], [B_init])
model = helper.make_model(graph)
datasets = DATASETS('onnxrt_qlinearops')
dataset = datasets['dummy'](shape=(1, 1, 5, 5), label=True)
dataloader = DATALOADERS['onnxrt_qlinearops'](dataset)
return model, dataloader
def test_tensorflow_dummy_v2(self):
datasets = DATASETS('tensorflow')
# test with label
dataset = datasets['dummy_v2'](\
input_shape=(256, 256, 3), label_shape=(1,))
data_loader = DATALOADERS['tensorflow'](dataset)
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data[0].shape, (1, 256, 256, 3))
self.assertEqual(data[1].shape, (1, 1))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data[0].shape, (2, 256, 256, 3))
self.assertEqual(data[1].shape, (2, 1))
# test without label
dataset = datasets['dummy_v2'](input_shape=(256, 256, 3))
data_loader = DATALOADERS['tensorflow'](dataset)
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (1, 256, 256, 3))
# dynamic batching
data_loader.batch(batch_size=2, last_batch='rollover')
iterator = iter(data_loader)
data = next(iterator)
self.assertEqual(data.shape, (2, 256, 256, 3))
with self.assertRaises(AssertionError):
dataset = datasets['dummy_v2'](\
input_shape=(256, 256, 3), low=[1., 0.])
with self.assertRaises(AssertionError):
dataset = datasets['dummy_v2'](\
input_shape=(256, 256, 3), high=[128., 127.])
with self.assertRaises(AssertionError):
dataset = datasets['dummy_v2'](\
input_shape=(256, 256, 3), dtype=['float32', 'int8'])
def testLabelBalanceCOCORecord(self):
from PIL import Image
tf.compat.v1.disable_eager_execution()
random_array = np.random.random_sample([100, 100, 3]) * 255
random_array = random_array.astype(np.uint8)
im = Image.fromarray(random_array)
im.save('test.jpeg')
image = tf.compat.v1.gfile.FastGFile('test.jpeg', 'rb').read()
source_id = '000000397133.jpg'.encode('utf-8')
label = 'person'.encode('utf-8')
example1 = tf.train.Example(features=tf.train.Features(
feature={
'image/encoded':
tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),
'image/object/class/text':
tf.train.Feature(bytes_list=tf.train.BytesList(value=[label])),
'image/source_id':
tf.train.Feature(bytes_list=tf.train.BytesList(
value=[source_id])),
'image/object/bbox/xmin':
tf.train.Feature(float_list=tf.train.FloatList(value=[10])),
'image/object/bbox/ymin':
tf.train.Feature(float_list=tf.train.FloatList(value=[10])),
'image/object/bbox/xmax':
tf.train.Feature(float_list=tf.train.FloatList(value=[100])),
'image/object/bbox/ymax':
tf.train.Feature(float_list=tf.train.FloatList(value=[100])),
}))
example2 = tf.train.Example(features=tf.train.Features(
feature={
'image/encoded':
tf.train.Feature(bytes_list=tf.train.BytesList(value=[image])),
'image/object/class/text':
tf.train.Feature(bytes_list=tf.train.BytesList(value=[label])),
'image/source_id':
tf.train.Feature(bytes_list=tf.train.BytesList(
value=[source_id])),
'image/object/bbox/xmin':
tf.train.Feature(float_list=tf.train.FloatList(
value=[10, 20])),
'image/object/bbox/ymin':
tf.train.Feature(float_list=tf.train.FloatList(
value=[10, 20])),
'image/object/bbox/xmax':
tf.train.Feature(float_list=tf.train.FloatList(
value=[100, 200])),
'image/object/bbox/ymax':
tf.train.Feature(float_list=tf.train.FloatList(
value=[100, 200])),
}))
with tf.io.TFRecordWriter('test.record') as writer:
writer.write(example1.SerializeToString())
writer.write(example2.SerializeToString())
preprocesses = TRANSFORMS('tensorflow', 'preprocess')
preprocess = get_preprocess(preprocesses, {'ParseDecodeCoco': {}})
filters = FILTERS('tensorflow')
filter = filters['LabelBalanceCOCORecord'](2)
datasets = DATASETS('tensorflow')
dataset = datasets['COCORecord']('test.record', \
transform=preprocess, filter=filter)
dataloader = DATALOADERS['tensorflow'](dataset=dataset, batch_size=1)
for (inputs, labels) in dataloader:
self.assertEqual(inputs.shape, (1, 100, 100, 3))
self.assertEqual(labels[0].shape, (1, 2, 4))
dataset2 = create_dataset('tensorflow',
{'COCORecord': {
'root': 'test.record'
}}, {'ParseDecodeCoco': {}},
{'LabelBalance': {
'size': 2
}})
dataloader2 = DATALOADERS['tensorflow'](dataset=dataset2, batch_size=1)
for (inputs, labels) in dataloader2:
self.assertEqual(inputs.shape, (1, 100, 100, 3))
self.assertEqual(labels[0].shape, (1, 2, 4))
dataloader3 = create_dataloader('tensorflow', {'batch_size':1, 'dataset':{'COCORecord':{'root':'test.record'}},\
'filter':{'LabelBalance':{'size':2}}, 'transform':{'ParseDecodeCoco':{}}})
for (inputs, labels) in dataloader3:
self.assertEqual(inputs.shape, (1, 100, 100, 3))
self.assertEqual(labels[0].shape, (1, 2, 4))
os.remove('test.record')
os.remove('test.jpeg')
def testLabelBalanceCOCORaw(self):
random_array = np.random.random_sample([100,100,3]) * 255
random_array = random_array.astype(np.uint8)
im = Image.fromarray(random_array)
os.makedirs('val2017', exist_ok=True)
im.save('./val2017/test_0.jpg')
im.save('./val2017/test_1.jpg')
fake_dict = {
'info': {
'description': 'COCO 2017 Dataset',
'url': 'http://cocodataset.org',
'version': '1.0',
'year': 2017,
'contributor': 'COCO Consortium',
'date_created': '2017/09/01'
},
'licenses':{},
'images':[{
'file_name': 'test_0.jpg',
'height': 100,
'width': 100,
'id': 0
},
{
'file_name': 'test_1.jpg',
'height': 100,
'width': 100,
'id': 1
}],
'annotations':[{
'category_id': 18,
'id': 1767,
'iscrowd': 0,
'image_id': 0,
'bbox': [473.07, 395.93, 38.65, 28.67],
},
{
'category_id': 18,
'id': 1768,
'iscrowd': 0,
'image_id': 1,
'bbox': [473.07, 395.93, 38.65, 28.67],
},
{
'category_id': 18,
'id': 1768,
'iscrowd': 0,
'image_id': 1,
'bbox': [473.07, 395.93, 38.65, 28.67],
}],
'categories':[{
'supercategory': 'animal',
'id': 18,
'name': 'dog'
}]
}
fake_json = json.dumps(fake_dict)
os.makedirs('annotations', exist_ok=True)
with open('./annotations/instances_val2017.json', 'w') as f:
f.write(fake_json)
filters = FILTERS('onnxrt_qlinearops')
filter = filters['LabelBalanceCOCORaw'](1)
datasets = DATASETS('onnxrt_qlinearops')
dataset = datasets['COCORaw']('./', transform=None, filter=filter)
dataloader = DATALOADERS['onnxrt_qlinearops'](dataset=dataset, batch_size=1)
for (inputs, labels) in dataloader:
self.assertEqual(labels[0].shape[1], 1)
shutil.rmtree('annotations')
shutil.rmtree('val2017')
class TestAdaptorONNXRT(unittest.TestCase):
mb_v2_export_path = "mb_v2.onnx"
mb_v2_model = torchvision.models.mobilenet_v2()
rn50_export_path = "rn50.onnx"
rn50_model = torchvision.models.resnet50()
datasets = DATASETS('onnxrt_qlinearops')
cv_dataset = datasets['dummy'](shape=(100, 3, 224, 224), low=0., high=1., label=True)
cv_dataloader = DATALOADERS['onnxrt_qlinearops'](cv_dataset)
ir3_dataset = datasets['dummy'](shape=(10, 2048), low=0., high=1., label=True)
ir3_dataloader = DATALOADERS['onnxrt_qlinearops'](ir3_dataset)
@classmethod
def setUpClass(self):
build_static_yaml()
build_dynamic_yaml()
build_non_MSE_yaml()
build_benchmark_yaml()
export_onnx_model(self.mb_v2_model, self.mb_v2_export_path)
self.mb_v2_model = onnx.load(self.mb_v2_export_path)
export_onnx_model(self.rn50_model, self.rn50_export_path)
self.rn50_model = onnx.load(self.rn50_export_path)
self.ir3_model = build_ir3_model()
@classmethod
def tearDownClass(self):
os.remove("static.yaml")
os.remove("dynamic.yaml")
os.remove("non_MSE.yaml")
os.remove("benchmark.yaml")
os.remove(self.mb_v2_export_path)
os.remove(self.rn50_export_path)
shutil.rmtree("./saved", ignore_errors=True)
shutil.rmtree("runs", ignore_errors=True)
def test_adaptor(self):
framework_specific_info = {"device": "cpu",
"approach": "post_training_static_quant",
"random_seed": 1234,
"q_dataloader": None,
"backend": "qlinearops",
"workspace_path": './lpot_workspace/{}/{}/'.format(
'onnxrt',
'imagenet')}
framework = "onnxrt_qlinearops"
adaptor = FRAMEWORKS[framework](framework_specific_info)
adaptor.inspect_tensor(self.rn50_model, self.cv_dataloader, inspect_type='activation')
adaptor.inspect_tensor(self.rn50_model, self.cv_dataloader, inspect_type='activation', save_to_disk=True)
adaptor.inspect_tensor(self.rn50_model, self.cv_dataloader, inspect_type='weight')
adaptor.inspect_tensor(self.rn50_model, self.cv_dataloader, inspect_type='all')
adaptor.inspect_tensor(self.rn50_model, self.cv_dataloader, ["Conv_0"], inspect_type='activation')
def test_set_tensor(self):
quantizer = Quantization("static.yaml")
quantizer.calib_dataloader = self.cv_dataloader
quantizer.eval_dataloader = self.cv_dataloader
quantizer.model = common.Model(self.mb_v2_model)
q_model = quantizer()
framework_specific_info = {"device": "cpu",
"approach": "post_training_static_quant",
"random_seed": 1234,
"q_dataloader": None,
"backend": "qlinearops",
"workspace_path": './lpot_workspace/{}/{}/'.format(
'onnxrt',
'imagenet')}
framework = "onnxrt_qlinearops"
adaptor = FRAMEWORKS[framework](framework_specific_info)
q_config = {'fused Conv_0': {'weight': {'granularity': 'per_channel', 'dtype': onnx_proto.TensorProto.INT8}}}
adaptor.q_config = q_config
adaptor.set_tensor(q_model, {'ConvBnFusion_W_features.0.0.weight': np.random.random([32, 3, 3, 3])})
adaptor.set_tensor(q_model, {'ConvBnFusion_BN_B_features.0.1.bias': np.random.random([32])})
def test_adaptor(self):
for fake_yaml in ["static.yaml", "dynamic.yaml"]:
quantizer = Quantization(fake_yaml)
quantizer.calib_dataloader = self.cv_dataloader
quantizer.eval_dataloader = self.cv_dataloader
quantizer.model = common.Model(self.rn50_model)
q_model = quantizer()
eval_func(q_model)
for fake_yaml in ["non_MSE.yaml"]:
quantizer = Quantization(fake_yaml)
quantizer.calib_dataloader = self.cv_dataloader
quantizer.eval_dataloader = self.cv_dataloader
quantizer.model = common.Model(self.mb_v2_model)
q_model = quantizer()
eval_func(q_model)
for fake_yaml in ["static.yaml"]:
quantizer = Quantization(fake_yaml)
quantizer.calib_dataloader = self.ir3_dataloader
quantizer.eval_dataloader = self.ir3_dataloader
quantizer.model = common.Model(self.ir3_model)
q_model = quantizer()
for mode in ["performance", "accuracy"]:
fake_yaml = "benchmark.yaml"
evaluator = Benchmark(fake_yaml)
evaluator.b_dataloader = self.cv_dataloader
evaluator.model = common.Model(self.rn50_model)
evaluator(mode)
def main(args=None):
tf.logging.set_verbosity(tf.logging.INFO)
if not tf.gfile.Exists(FLAGS.output_dir):
tf.gfile.MkDir(FLAGS.output_dir)
with tf.Session() as sess:
if FLAGS.input_model.rsplit('.', 1)[-1] == 'ckpt':
style_img_ph = tf.placeholder(tf.float32,
shape=[None, 256, 256, 3],
name='style_input')
content_img_ph = tf.placeholder(tf.float32,
shape=[None, 256, 256, 3],
name='content_input')
# import meta_graph
meta_data_path = FLAGS.input_model + '.meta'
saver = tf.train.import_meta_graph(meta_data_path,
clear_devices=True)
sess.run(tf.global_variables_initializer())
saver.restore(sess, FLAGS.input_model)
graph_def = sess.graph.as_graph_def()
replace_style = 'style_image_processing/ResizeBilinear_2'
replace_content = 'batch_processing/batch'
for node in graph_def.node:
for idx, input_name in enumerate(node.input):
# replace style input and content input nodes to placeholder
if replace_content == input_name:
node.input[idx] = 'content_input'
if replace_style == input_name:
node.input[idx] = 'style_input'
if FLAGS.tune:
_parse_ckpt_bn_input(graph_def)
output_name = 'transformer/expand/conv3/conv/Sigmoid'
frozen_graph = tf.graph_util.convert_variables_to_constants(
sess, graph_def, [output_name])
# use frozen pb instead
elif FLAGS.input_model.rsplit('.', 1)[-1] == 'pb':
with open(FLAGS.input_model, 'rb') as f:
frozen_graph = tf.GraphDef()
frozen_graph.ParseFromString(f.read())
else:
print("not supported model format")
exit(-1)
if FLAGS.tune:
with tf.Graph().as_default() as graph:
tf.import_graph_def(frozen_graph, name='')
quantizer = Quantization(FLAGS.config)
quantizer.model = graph
quantized_model = quantizer()
quantized_model.save(FLAGS.output_model)
frozen_graph = quantized_model.graph_def
# validate the quantized model here
with tf.Graph().as_default(), tf.Session() as sess:
if FLAGS.tune:
# create dataloader using default style_transfer dataset
# generate stylized images
dataset = DATASETS('tensorflow')['style_transfer']( \
FLAGS.content_images_paths.strip(),
FLAGS.style_images_paths.strip(),
crop_ratio=0.2,
resize_shape=(256, 256))
else:
dataset = DATASETS('tensorflow')['dummy']( \
shape=[(200, 256, 256, 3), (200, 256, 256, 3)], label=True)
dataloader = DATALOADERS['tensorflow'](dataset=dataset,
batch_size=FLAGS.batch_size)
tf.import_graph_def(frozen_graph, name='')
style_transfer(sess, dataloader, FLAGS.precision)