Hi everyone. I'm new to TVM and I'm trying to tune a simple onnx model using `auto_scheduler.TaskScheduler` on my CPU. The tuning process doesn't finish even after several hours and doesn't either return me any logs. The model has a simple structure so I assume tuning it should require a lot of computation. The code for the model:
class Net(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 6, 5) self.pool = nn.MaxPool2d(2, 2) self.conv2 = nn.Conv2d(6, 16, 5) self.fc1 = nn.Linear(16 * 5 * 5, 120) self.fc2 = nn.Linear(120, 84) self.fc3 = nn.Linear(84, 10) def forward(self, x): x = self.pool(F.relu(self.conv1(x))) x = self.pool(F.relu(self.conv2(x))) x = torch.flatten(x, 1) # flatten all dimensions except batch x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x net = Net() net.load_state_dict(torch.load(PATH)) batch_size = 4 import torch.onnx net.train(False) x = torch.randn(batch_size, 3, 32, 32, requires_grad=True) torch_out = torch.onnx._export(net, # model being run x, # model input (or a tuple for multiple inputs) "testModel1.onnx", # where to save the model (can be a file or file-like object) export_params=True) The code for the tuning : if __name__ == '__main__': target = "llvm" network = "testModel1" layout = "NCHW" input_name = "input.1" # shape_dict = {input_name: images.shape} shape_dict = {input_name: torch.Size([batch_size, 3, 32, 32])} mod, params = relay.frontend.from_onnx(onnx_model, shape_dict) log_file = "%s-%s-B%d-%s.json" % (network, layout, batch_size, target) print(log_file) print("Extract tasks...") tasks, task_weights = auto_scheduler.extract_tasks(mod["main"], params, target) for idx, task in enumerate(tasks): print("========== Task %d (workload key: %s) ==========" % (idx, task.workload_key)) print(task.compute_dag) def run_tuning(): print("Begin tuning...") tuner = auto_scheduler.TaskScheduler(tasks, task_weights) tune_option = auto_scheduler.TuningOptions( num_measure_trials=1, # change this to 20000 to achieve the best performance runner=auto_scheduler.LocalRunner(repeat=1, enable_cpu_cache_flush=True), measure_callbacks=[auto_scheduler.RecordToFile(log_file)], ) tuner.tune(tune_option) run_tuning() # Compile with the history best print("Compile...") with auto_scheduler.ApplyHistoryBest(log_file): with tvm.transform.PassContext(opt_level=3, config={"relay.backend.use_auto_scheduler": True}): lib = relay.build(mod, target=target, params=params) # Create graph executor dev = tvm.device(target, 0) module = graph_executor.GraphModule(lib["default"](dev)) data_tvm = tvm.nd.array((np.random.uniform(size=torch.Size([batch_size, 3, 32, 32]))).astype("float32")) module.set_input("input.1", data_tvm) # Evaluate print("Evaluate inference time cost...") print(module.benchmark(dev, repeat=3, min_repeat_ms=500)) my system: -Processor: Inter(R) Core(TM) i5 CPU 750 @ 2.67GHz 2.67 GHz -Installed RAM: 16.0 GB -System type: 64-bit operating system, x64-based processor --- [Visit Topic](https://discuss.tvm.apache.org/t/tuning-doesnt-finish-even-after-several-hours/12072/1) to respond. You are receiving this because you enabled mailing list mode. To unsubscribe from these emails, [click here](https://discuss.tvm.apache.org/email/unsubscribe/89118bd0cd8eb044c169796ac8d47cad87f450b3730f342ec34cccabb92ae4be).