From 4c77aef489552851e6b22b65363e7ae45b27a2f1 Mon Sep 17 00:00:00 2001 From: Cody Hao Yu Date: Sat, 14 Dec 2019 00:30:14 +0000 Subject: [PATCH] fix empty config caused KeyError --- python/tvm/autotvm/record.py | 8 +++++++- tests/python/unittest/test_autotvm_record.py | 7 +++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/python/tvm/autotvm/record.py b/python/tvm/autotvm/record.py index 14efb7bd9239..fbf4a08f7b0c 100644 --- a/python/tvm/autotvm/record.py +++ b/python/tvm/autotvm/record.py @@ -183,7 +183,13 @@ def load_from_file(filename): """ for row in open(filename): if row and not row.startswith('#'): - yield decode(row) + inp, res = decode(row) + # Avoid loading the record with an empty config. The TOPI schedule with no entities + # will result in an empty entity map (e.g., depthwise_conv2d_nchw on x86). + # Using an empty config will cause problems when applying alter op like NCHW to NCHWc. + if not inp.config._entity_map: + continue + yield (inp, res) def split_workload(in_file, clean=True): diff --git a/tests/python/unittest/test_autotvm_record.py b/tests/python/unittest/test_autotvm_record.py index f83dd7c74d0b..0839ad9b68cf 100644 --- a/tests/python/unittest/test_autotvm_record.py +++ b/tests/python/unittest/test_autotvm_record.py @@ -52,9 +52,16 @@ def test_file_io(): inputs = [MeasureInput(target, tsk, tsk.config_space.get(i)) for i in range(0, 10)] results = [MeasureResult((i, ), 0, 0, 0) for i in range(0, 10)] + invalid_inp = MeasureInput(target, tsk, tsk.config_space.get(10)) + invalid_res = MeasureResult((10, ), 0, 0, 0) + + # Erase the entity map to test if it will be ignored when loading back. + invalid_inp.config._entity_map = {} + with open(file_path, "w") as fo: cb = autotvm.callback.log_to_file(fo) cb(None, inputs, results) + cb(None, [invalid_inp], [invalid_res]) ref = zip(inputs, results) for x, y in zip(ref, autotvm.record.load_from_file(file_path)):