multiprocess

handcrafted_solution.py

@@ -239,4 +239,4 @@ def predict(entry, visualize=False) -> Tuple[np.ndarray, List[int]]:
         from viz3d import plot_estimate_and_gt
         plot_estimate_and_gt(all_3d_vertices_clean, connections_3d_clean, good_entry['wf_vertices'],
                             good_entry['wf_edges'])
-    return  all_3d_vertices_clean, connections_3d_clean, [0 for i in range(len(connections_3d_clean))]
+    return good_entry['__key__'], all_3d_vertices_clean, connections_3d_clean, [0 for i in range(len(connections_3d_clean))]

script.py

@@ -143,18 +143,24 @@ if __name__ == "__main__":
     dataset = hoho.get_dataset(decode=None, split='all', dataset_type='webdataset')
     print('------------ Now you can do your solution ---------------')
     solution = []
-    for i, sample in enumerate(tqdm(dataset)):
-        pred_vertices, pred_edges, semantics = predict(sample, visualize=False)
-        solution.append({
-                        '__key__': sample['__key__'], 
-                        'wf_vertices': pred_vertices.tolist(),
-                        'wf_edges': pred_edges,
-                        'edge_semantics': semantics,
-                    })
-        if i % 100 == 0:
-            # incrementally save the results in case we run out of time
-            print(f"Processed {i} samples")
-            # save_submission(solution, Path(params['output_path']) / "submission.parquet")
+    from concurrent.futures import ProcessPoolExecutor
+    with ProcessPoolExecutor(max_workers=8) as pool:
+        results = []
+        for i, sample in enumerate(tqdm(dataset)):
+            results.append(pool.submit(predict, sample, visualize=False))
+        
+        for i, result in enumerate(tqdm(results)):
+            key = pred_vertices, pred_edges, semantics = result.result()
+            solution.append({
+                            '__key__': key,
+                            'wf_vertices': pred_vertices.tolist(),
+                            'wf_edges': pred_edges,
+                            'edge_semantics': semantics,
+                        })
+            if i % 100 == 0:
+                # incrementally save the results in case we run out of time
+                print(f"Processed {i} samples")
+                # save_submission(solution, Path(params['output_path']) / "submission.parquet")
     print('------------ Saving results ---------------')
     save_submission(solution, Path(params['output_path']) / "submission.parquet")
     print("------------ Done ------------ ")