yuchaoli commented on a change in pull request #8065:
URL: https://github.com/apache/tvm/pull/8065#discussion_r634060352



##########
File path: tutorials/auto_scheduler/tune_network_x86.py
##########
@@ -138,7 +140,9 @@ def get_network(name, batch_size, layout="NHWC", 
dtype="float32", use_sparse=Fal
     if use_sparse:
         from tvm.topi.sparse.utils import convert_model_dense_to_sparse
 
-        mod, params = convert_model_dense_to_sparse(mod, params, 
random_params=True)
+        mod, params = convert_model_dense_to_sparse(
+            mod, params, bs_r=sparse_block_size[0], bs_c=sparse_block_size[1], 
random_params=True

Review comment:
       Done.

##########
File path: tutorials/auto_scheduler/tune_network_x86.py
##########
@@ -205,7 +215,33 @@ def run_tuning():
         measure_callbacks=[auto_scheduler.RecordToFile(log_file)],
     )
 
-    tuner.tune(tune_option)
+    if use_sparse:
+        from tvm.topi.sparse.utils import (
+            sparse_conv2d_meet_condition_func,
+            sparse_conv2d_apply_func,
+            sparse_dense_meet_condition_func,
+            sparse_dense_apply_func,
+        )
+
+        search_policy = [
+            auto_scheduler.SketchPolicy(
+                task,
+                program_cost_model=auto_scheduler.XGBModel(),
+                init_search_callbacks=[
+                    auto_scheduler.PreloadCustomSketchRule(
+                        sparse_conv2d_meet_condition_func, 
sparse_conv2d_apply_func, "SparseConv2D"
+                    ),
+                    auto_scheduler.PreloadCustomSketchRule(
+                        sparse_dense_meet_condition_func, 
sparse_dense_apply_func, "SparseDense"
+                    ),
+                ],
+            )
+            for task in tasks
+        ]

Review comment:
       Done.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to