{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":11244260,"sourceType":"datasetVersion","datasetId":6993770}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Speed up inference\n\nFast inference was an important element of the similar competition\n[CZII](https://www.kaggle.com/competitions/czii-cryo-et-object-identification/overview). \n\n\n**Preprocess on GPU:**\n\nPreprocessing on CPU (resize and normalize) can take significant time comparable to model inference, but become negligible if you do it on GPU.\n\nFor 20 training tomo_ids with one model, the inference time is:\n\n```text\n240 sec CPU preprocessing\n 65 sec GPU\n```\n\n\n**Use amp with T4×2:**\n\nT4×2 is faster than P100×1 for float32 only marginary, but\nautomatic mixed precision (amp) speeds up by a factor of 2 on T4, while not on P100.\n\n```text\n160 sec  P100 float32 num_workers=2\n135 sec  T4×2 float32 num_workers=1\n 65 sec  T4×2 amp     num_workers=1\n```\n\n(20 trian tomo_ids, one model)\n\n\nI am also interested in:\n\n* TensorRT\n  \nbut not in this notebook (yet?).\n\n**Scores**\n\n```text\nLocal CV:\n  0.8164 ± 0.0355\n  5 folds [0.7612 0.8076 0.8449 0.8495 0.8186]\n\nPublic LB:\n  Fold 0       0.583   66 min\n  5-fold mean  0.665  227 min\n```\n\nVersion: 2\n\nFix wrong y_pred in predict(), thanks to\n[sacuscreed's comment](https://www.kaggle.com/code/junkoda/speed-up-inference/comments#3170678),\nbut the Public LB 0.665 remains the same.\n\n**Version: 3**\nImplement Test-Time Augmentation (TTA).\n\n```text\nLocal CV (4 rot90 TTA)\n  0.8536 ± 0.0329\n  5 folds [0.8311 0.8171 0.8824 0.8933 0.8442]\n\nPublic LB:\n  model0 x 4TTA  0.668  184 min\n```\n\nI thought the 5-fold average is better than 1 model x 4 TTA because of the models have more varieties,\nso I do not conclude that K-TTA is better than K-model average in general, but might be comparable.\nYou can also use TTA for Local CV, which could make the evaluation a little more robust.","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport time\nimport argparse\nfrom pathlib import Path\n\nfrom tqdm.auto import tqdm\nfrom torch.utils.data import DataLoader, default_collate\nfrom PIL import Image\n\nimport torch\nimport torch.nn as nn\nimport torchvision.transforms.functional as TTF\nimport timm\nimport yaml\n\nimport multiprocessing as mp\nfrom queue import Empty\n\n\nINPUT_PATH = Path('/kaggle/input/byu-locating-bacterial-flagellar-motors-2025')\nMODEL_PATH = Path('/kaggle/input/bacterial-public/weights/object/baseline')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:04:48.191442Z","iopub.execute_input":"2025-04-07T08:04:48.191832Z","iopub.status.idle":"2025-04-07T08:05:01.882409Z","shell.execute_reply.started":"2025-04-07T08:04:48.191801Z","shell.execute_reply":"2025-04-07T08:05:01.881683Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data","metadata":{}},{"cell_type":"code","source":"def get_tomos(input_path: Path,\n              data_type: str,\n              *, n=None) -> list[Path]:\n    \"\"\"\n    tomo..dachi...?\n    Args\n      input_path (Path): Kaggle input directory\n      data_type (str):   train or test\n      n (Optional[int]): Use only first n train tomos (not applicable to test)\n    \"\"\"\n    data_path = input_path / data_type\n    tomo_paths = sorted(data_path.glob('*'))\n\n    if (n is not None) and (data_type == 'train'):\n        tomo_paths = tomo_paths[:n]\n    \n    return tomo_paths\n    \n\ndef preprocess(img: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n    Resize and normalize\n\n    Arg:\n      img (Tensor[uint8]):   (batch_size, C, H, W)\n\n    Returns:\n      img (Tensor[float32]): (batch_size, C, size, size); size = 640\n    \"\"\"\n    size = (640, 640)\n\n    img = img.to(dtype=torch.float32)\n    img = TTF.resize(img, size)  # (batch_size, C, size, size)\n\n    batch_size, nch, h, w = img.shape\n    q = torch.Tensor([0.05, 0.95]).to(img.device)\n    x_min, x_max = torch.quantile(img.view(batch_size, nch * h * w), q, dim=1)\n    x_min = x_min.view(batch_size, 1, 1, 1)\n    x_max = x_max.view(batch_size, 1, 1, 1)\n\n    img = (img - x_min) / (x_max - x_min)\n    img = torch.clamp(img, 0, 1)\n\n    return img\n\n\nclass Dataset(torch.utils.data.Dataset):\n    \"\"\"\n    dataset = Dataset(tomo_path)\n    \n    Args:\n      tomo_path (Path): directory including jpg images    \n    \"\"\"\n    def __init__(self, tomo_path: Path):\n        self.filenames = sorted(tomo_path.glob('*'))  # list[Path]\n\n    def __len__(self) -> int:\n        return len(self.filenames)\n\n    def __getitem__(self, i: int) -> dict:\n        filename = self.filenames[i]  # Path\n        filebase = filename.stem\n        assert filebase[:6] == 'slice_'\n        slice_number = int(filebase[6:])  # slice_0000 -> int(0000)\n\n        # Load, resize and normalize image\n        img = Image.open(filename)\n        W, H = img.size\n        img = np.expand_dims(np.array(img), axis=0)  # array[uint8] (1, H, W)\n\n        ret = {'img': img,\n               'slice_number': slice_number,\n               'shape': np.array((H, W), dtype=int),  # original shape H, W\n        }\n\n        return ret\n\n    def loader(self, batch_size: int, num_workers: int):\n        loader = DataLoader(self, batch_size=batch_size, num_workers=num_workers)\n        return loader\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:01.883395Z","iopub.execute_input":"2025-04-07T08:05:01.883619Z","iopub.status.idle":"2025-04-07T08:05:01.893176Z","shell.execute_reply.started":"2025-04-07T08:05:01.883598Z","shell.execute_reply":"2025-04-07T08:05:01.892154Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Model\n\nThe model is a simple 2D object detection using convnext_tiny.\n\nThe main theme of this notebook is inference,\nbut if you are intersted in this model,\nsee github:\n\nhttps://github.com/junkoda/kaggle_bacterial_public\n\nI followed sharifi76's great notebook:\n\nhttps://www.kaggle.com/code/sharifi76/eda-visualization-yolov8\n\nusing only ±4 positive slices around the label for training.\n\nUsually using all data is best in terms of score, but you can speedup training significantly by reducing negative samples with little loss in score. I haven't tried yet but you can probably improve score by adding negative samples and adjust threshold for no moters (and also with more augmentations).\n\n\n* Model input is 640x640 image\n* Pretrained model outputs 20x20 embedding vectors\n* Predict 20x20\n    + probabilty including target\n    + offset within the coarse grid\n* Object detection models also predict sizes of bounding boxes, but this model only predicts the center points.\n","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport timm\n\n\nclass Model(nn.Module):\n    def __init__(self, cfg_model: dict, *, pretrained=True, verbose=True):\n        super().__init__()\n\n        # Timm encoder\n        name = cfg_model['encoder']\n        in_channels = 1\n        out_channels = 1\n\n        self.encoder = timm.create_model(name,\n                                         in_chans=in_channels,\n                                         features_only=True,\n                                         pretrained=pretrained)\n        encoder_channels = self.encoder.feature_info.channels()\n\n        self.segmentation_head = nn.Conv2d(encoder_channels[-1], out_channels,\n                                           kernel_size=3, padding=1)\n\n        self.regression_head = nn.Conv2d(encoder_channels[-1], out_channels=2,\n                                         kernel_size=3, padding=1)\n\n        self.criterion_seg = nn.BCEWithLogitsLoss()\n        self.criterion_reg = nn.MSELoss()\n\n        if verbose:\n            print(name)\n\n    def forward(self, img: torch.Tensor):\n        \"\"\"\n        img (Tensor): (batch_size, 1, H, W)\n\n        h, w = H // 32, W // 32\n        \"\"\"\n        features = self.encoder(img)\n        out = features[-1]                    # (batch_size, embed_dim, h, w)\n        y_pred = self.segmentation_head(out)  # (batch_size, 1, h, w)\n        t_pred = self.regression_head(out)    # (batch_size, 2, h, w)\n\n        return y_pred, t_pred\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:01.894799Z","iopub.execute_input":"2025-04-07T08:05:01.895032Z","iopub.status.idle":"2025-04-07T08:05:01.929428Z","shell.execute_reply.started":"2025-04-07T08:05:01.895013Z","shell.execute_reply":"2025-04-07T08:05:01.928649Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Predict","metadata":{}},{"cell_type":"code","source":"def predict(tomo_path: Path,\n            models: list[nn.Module],\n            cfg: dict) -> dict:\n    \"\"\"\n    Predict moter coordinate for one tomo_id\n    At most one moter per tomo_id in test\n    \n    Args:\n      model (nn.Module): Pytorch model\n      dataset (Dataset): for one tomo_id\n    \"\"\"\n    assert len(models) > 0\n    tomo_id = tomo_path.name\n    batch_size = cfg['batch_size']\n    num_workers = cfg['num_workers']\n    use_amp = cfg['use_amp']\n    preprocess_device = cfg['preprocess_device']\n    assert preprocess_device in ['cuda', 'cpu']\n\n    dataset = Dataset(tomo_path)\n    loader = dataset.loader(batch_size=batch_size, num_workers=num_workers)\n    \n    device = next(models[0].parameters()).device\n\n    K = len(cfg['tta'])   # 4 for 4 rot90, 1 for no TTA\n    tta = TTA(cfg['tta'], device=device)\n\n    # Loop over all slices in one tomo_id\n    best = (0, None)\n    for d in loader:\n        # Input image (batch_size, 1, H, W)\n        if preprocess_device == 'cuda':\n            img = d['img'].to(device) \n            img = preprocess(img)\n        elif preprocess_device == 'cpu':\n            img = preprocess(d['img'])\n            img = img.to(device)  # input image (batch_size, 1, H, W)\n        else:\n            raise ValueError(preprocess_device)\n\n        # TTA: batch_size -> K * batch_size\n        img = tta.expand(img)   # (K * batch_size, 1, H, W)\n        \n        y_pred_sum, t_pred_sum = None, None\n        for model in models:\n            with torch.no_grad():\n                with torch.amp.autocast(device_type='cuda',\n                                        enabled=use_amp,\n                                        dtype=torch.float16):\n                    y_pred, t_pred = model(img) \n\n            y_pred = y_pred.sigmoid()  # (batch_size, 1, h, w)\n\n            if y_pred_sum is None:\n                y_pred_sum = y_pred\n                t_pred_sum = t_pred\n            else:\n                y_pred_sum += y_pred\n                t_pred_sum += t_pred\n\n        # TTA: Reduce K * batch_size -> batch_size\n        y_pred_sum, t_pred_sum = tta.reduce(y_pred_sum, t_pred_sum)  # (batch_size, 1, h, w)  \n        \n        y_pred_max = y_pred_sum.max().item() / len(models)\n        del y_pred, t_pred\n\n        # Keep most probable coordinate\n        if y_pred_max > best[0]:\n            bs, _, h, w = y_pred_sum.shape\n        \n            argmax = torch.unravel_index(y_pred_sum.argmax(), y_pred_sum.shape)  # b, ch, iy, ix\n            i, _, iy, ix = [t.item() for t in argmax]\n            slice_number = d['slice_number'][i].item()\n            offset = t_pred_sum[i, :, iy, ix].cpu().numpy() / len(models)  # (2, )\n\n            # Compute coodinate in original pixels\n            H, W = d['shape'][i].numpy()    # Original image size\n            x = (ix + offset[0]) * (W / w)\n            y = (iy + offset[1]) * (H / h)\n            \n            best = (y_pred_max, slice_number, y, x)\n\n    assert best[1] is not None\n\n    # Return prediction\n    n_slices = len(dataset.filenames)\n    pred = {'tomo_id': tomo_id,\n            'n_slices': n_slices,\n            'y_pred': best[0],\n            'zyx': best[1:]}\n    return pred\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:01.930486Z","iopub.execute_input":"2025-04-07T08:05:01.930684Z","iopub.status.idle":"2025-04-07T08:05:01.947194Z","shell.execute_reply.started":"2025-04-07T08:05:01.930665Z","shell.execute_reply":"2025-04-07T08:05:01.946422Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Test-Time Augmentation (TTA)","metadata":{}},{"cell_type":"code","source":"def create_rot_matrices(ks, inverse=False) -> torch.Tensor:\n    \"\"\"\n    Create rot90 rotation matrices\n\n    Arg:\n      ks (list[int]): list of rot90 indices k=0,1,2,3\n      inverse (bool): If True kth rotation is rot90(-k)\n\n    Returns:\n      rots (Tensor[float32]): (K, 1, 2, 2) where K = len(ks)\n    \"\"\"\n    # Create four rot90 matrices\n    a = np.array([(1, 0), (0, 1)], dtype=np.float32)         # identity\n    if inverse:\n        r90 = np.array([(0, 1), (-1, 0)], dtype=np.float32)  # rotate -90 degrees\n    else:\n        r90 = np.array([(0, -1), (1, 0)], dtype=np.float32)  # rotate  90 degrees\n\n    rots = [a, ]\n    for k in range(4):\n        a = a @ r90\n        rots.append(a)\n\n    rots = [rots[k] for k in ks]\n\n    rots = np.stack(rots, axis=0)        # (K, 2, 2)\n    rots = np.expand_dims(rots, axis=1)  # (K, 1, 2, 2)\n    rots = torch.from_numpy(rots)\n\n    return rots\n\n\nclass TTA:\n    \"\"\"\n    tta = TTA(ks, dtype, device)\n\n    Args:\n      ks (list[int]): list of rot90 indices k=0,1,2,3\n      dtype (torch.dtype): torch.float16 if\n    \"\"\"\n    def __init__(self,\n                 ks: list[int],\n                 device: torch.device):\n        self.ks, self.no_tta = self._check_ks(ks)\n        self.rots = create_rot_matrices(ks, inverse=True).to(device=device)  # (K, 1, 2, 2)\n        self.no_tta = (list(ks) == [0, ])\n\n    @staticmethod\n    def _check_ks(ks):\n        for k in ks:\n            assert 0 <= k < 4\n        no_tta = (list(ks) == [0, ])\n\n        return ks, no_tta\n\n    def __repr__(self):\n        return 'TTA(%s)' % ','.join(['%d' % k for k in self.ks])\n\n    def expand(self, img: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Expand input image to K rot90 images\n\n        Args:\n          img (Tensor): (B, C, H, W)\n\n        Returns:\n          img_k (Tensor): (K * B, C, H, W)\n        \"\"\"\n        if self.no_tta:\n            return img\n\n        stack = []\n        for k in self.ks:\n            img_rot = torch.rot90(img, k, dims=[2, 3])\n            stack.append(img_rot)\n\n        return torch.cat(stack, dim=0)\n\n    def reduce(self,\n               img_k: torch.Tensor,\n               offset_k: torch.Tensor) -> tuple[torch.Tensor]:\n        \"\"\"\n        Average k patterns of rot90 predictions\n\n        Args:\n          img_k    (Tensor): (K * B, C, H, W)\n          offset_k (Tensor): (K * B, 2, H, W)\n\n        Returns: img, offset\n          img     (Tensor): (B, C, H, W)\n          offset  (Tensor): (B, C, H, W)\n        \"\"\"\n        if self.no_tta:\n            return img_k, offset_k\n\n        img = self.rotate_back_image(img_k, average=True)\n        offset = self.rotate_back_offset(offset_k, average=True)\n\n        return img, offset\n   \n    def rotate_back_image(self, img_k: torch.Tensor, *, average=True) -> torch.Tensor:\n        \"\"\"\n        Rotate back K TTA images (scaler fields)\n\n        Arg:\n          img_k (Tensor): (K * B, C, H, W)\n        \"\"\"\n        K = len(self.ks)\n        BK, C, H, W = img_k.shape\n        assert BK % K == 0\n        B = BK // K\n\n        img_k = img_k.view(K, B, C, H, W)\n\n        # Rotate back\n        if average:\n            # Reduce K TTA patterns by average\n            avg = torch.zeros((B, C, H, W), dtype=torch.float32, device=img_k.device)\n            for i, k in enumerate(self.ks):\n                avg += torch.rot90(img_k[i], 4 - k, dims=(2, 3))\n            avg /= K  # (B, C, H, W)\n        else:\n            out = []\n            for i, k in enumerate(self.ks):\n                out.append(torch.rot90(img_k[i], 4 - k, dims=(2, 3)))\n            return torch.cat(out, dim=0)  # (BK, C, H, W)\n\n        return avg\n\n    def rotate_back_offset(self,\n                           offset_k: torch.Tensor, *,\n                           average=True) -> torch.Tensor:\n        \"\"\"\n        Rotate back K TTA offsets (vector fields)\n\n        Arg:\n          offset (Tensor): (4 * B, 2, h, w)\n\n          offset[k] will be rotated by rot90(, k)\n          offset[:, :, 0] is dx and offset[:, :, 1] is dy from pixel cell top-left in [0, 1]\n        \"\"\"\n        K = len(self.ks)\n        BK, ndim, H, W = offset_k.shape\n        assert BK % K == 0\n        assert ndim == 2\n        B = BK // K\n\n        offset = offset_k.view(K, B, 2, H, W)\n        dtype = offset.dtype\n\n        # Expand rotation matrix to B batch\n        a = self.rots.expand(K, B, 2, 2).reshape(K * B, 2, 2).to(dtype=dtype)\n\n        # Offset vectors from pixel cell center\n        offset = (offset - 0.5).permute(0, 1, 3, 4, 2).reshape(K * B, H * W, 2)\n        \n        # Apply rotation matrix\n        offset_rot = torch.bmm(offset, a)\n        offset_rot = offset_rot.reshape(K, B, H, W, 2).permute(0, 1, 4, 2, 3)\n\n        # Back to offset from pixel cell top-left corner\n        offset_rot += 0.5  # (4, B, 2, H, W)\n\n        # Rotate pixels\n        if average:\n            # Reduce K TTA patterns by average\n            avg = torch.zeros((B, 2, H, W), dtype=torch.float32, device=offset.device)\n            for i, k in enumerate(self.ks):\n                avg += torch.rot90(offset_rot[i], 4 - k, dims=(2, 3))\n            avg /= K  # (B, 2, H, W)\n        else:\n            out = []\n            for i, k in enumerate(self.ks):\n                out.append(torch.rot90(offset_rot[i], 4 - k, dims=(2, 3)))\n            return torch.cat(out, dim=0)  # (BK, 2, H, W)\n        \n        return avg\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:01.948045Z","iopub.execute_input":"2025-04-07T08:05:01.948275Z","iopub.status.idle":"2025-04-07T08:05:01.974663Z","shell.execute_reply.started":"2025-04-07T08:05:01.948246Z","shell.execute_reply":"2025-04-07T08:05:01.974108Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submit","metadata":{}},{"cell_type":"code","source":"def create_submission(preds: list, th: float, ofilename: str) -> pd.DataFrame:\n    \"\"\"\n    Args:\n      preds (list[dict]): predictions\n      th (float): threshold between no or one moter\n      ofilename (str): submission.csv\n    \"\"\"\n    rows = []\n    count_positive = 0\n    for pred in preds:\n        if pred['y_pred'] < th:\n            zyx = (-1, -1, -1)\n        else:\n            count_positive += 1\n            zyx = pred['zyx']\n\n        row = {'tomo_id': pred['tomo_id'],\n               'Motor axis 0': zyx[0],\n               'Motor axis 1': zyx[1],\n               'Motor axis 2': zyx[2]}\n        rows.append(row)\n\n    submit = pd.DataFrame(rows)\n    submit.to_csv(ofilename, float_format='%.8e', index=False)\n\n    print('Submit %s: %d positives / %d tomo_ids' % (ofilename, count_positive, len(rows)))\n\n    return submit","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:01.975472Z","iopub.execute_input":"2025-04-07T08:05:01.975753Z","iopub.status.idle":"2025-04-07T08:05:01.998867Z","shell.execute_reply.started":"2025-04-07T08:05:01.975734Z","shell.execute_reply":"2025-04-07T08:05:01.998253Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main","metadata":{}},{"cell_type":"code","source":"def process_fn(process_id: int,\n               tomo_queue,\n               pred_queue,\n               cfg: dict):\n    \"\"\"\n    Prediction process for each GPU\n\n    Args:\n      process_id (int): 0 or 1 for 2 GPUs\n      tomo_queue (): tomo_ids\n      pred_queue (): output predictions\n      cfg (dict): config\n    \"\"\"\n    device = torch.device('cuda:%d' % process_id)\n\n    model_path = cfg['model_path']\n    folds = cfg['folds']\n\n    # Load models\n    with open(model_path / 'config.yml', 'r') as f:\n        cfg_model = yaml.safe_load(f)\n            \n    models = []\n    for ifold in folds:\n        model_filename = '%s/model%d.pytorch' % (model_path, ifold)\n        model = Model(cfg_model['model'], pretrained=False, verbose=False)\n        model.load_state_dict(torch.load(model_filename, weights_only=True))\n        model.to(device)\n        model.eval()\n        models.append(model)\n\n        if process_id == 0:\n            print('Load model', model_filename)\n        \n    # Loop over tomograms\n    while not tomo_queue.empty():\n        try:\n            tomo_path = tomo_queue.get(timeout=1)\n            pred = predict(tomo_path, models, cfg)\n            pred_queue.put(pred)\n\n        except Empty:\n            break\n            \n\n#\n# Main\n#\ntb = time.time()\n\n# Config\ncfg = {\n    'model_path': MODEL_PATH,\n    'folds': [0, ],          # [0,1,2,3,4] for all 5-fold mean\n    'batch_size': 16,\n    'num_workers': 1,\n    'use_amp': True,\n    'preprocess_device': 'cuda',\n    'tta': [0,1,2,3],              # rot90 ks, e.g., [0,] (no TTA) or [0,1,2,3] (all 4 rot90s)\n}\n\n\n#\n# List of tomograms\n#\ntomo_paths = get_tomos(INPUT_PATH, 'test')\n\nif len(tomo_paths) == 3:\n    # Some experiment when test is dummy (optional)\n    # tomo_paths = get_tomos(INPUT_PATH, 'train', n=20)\n    pass\n\nprint('Data %d' % len(tomo_paths))\nprint('TTA', cfg['tta'])\n\nmanager = mp.Manager()\ntomo_queue = manager.Queue()\npred_queue = manager.Queue()\nfor tomo_path in tomo_paths:\n    tomo_queue.put(tomo_path)\n\ntime.sleep(1)\nassert not tomo_queue.empty()\n\n\n#\n# Launch process\n#\nnum_processes = 2\ntb = time.time()\n\nworkers = [mp.Process(target=process_fn,\n                      args=(i, tomo_queue, pred_queue, cfg))\n           for i in range(num_processes)]\n\nfor w in workers:\n    w.start()\n\nfor w in workers:\n    w.join()\n\n\ndt = time.time() - tb\nprint('%.2f sec for %d tomos' % (dt, len(tomo_paths)))\n\n# queue to list\npreds = []\ntry:\n    while not pred_queue.empty():\n        preds.append(pred_queue.get(timeout=1))\nexcept Empty:\n    pass\n\nassert len(preds) == len(tomo_paths)\n\n\n#\n# Submit\n#\nth = 0.5\nofilename = 'submission.csv'\ncreate_submission(preds, th, 'submission.csv')\nprint(ofilename, 'written')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:01.999782Z","iopub.execute_input":"2025-04-07T08:05:02.000058Z","iopub.status.idle":"2025-04-07T08:05:36.671418Z","shell.execute_reply.started":"2025-04-07T08:05:02.000038Z","shell.execute_reply":"2025-04-07T08:05:36.670409Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! head -n 5 submission.csv\n! wc -l submission.csv","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-07T08:05:36.674320Z","iopub.execute_input":"2025-04-07T08:05:36.674596Z","iopub.status.idle":"2025-04-07T08:05:36.930581Z","shell.execute_reply.started":"2025-04-07T08:05:36.674570Z","shell.execute_reply":"2025-04-07T08:05:36.929822Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"> The first rule of program optimization: Don't do it.\n> \n> The second rule of program optimization: Don't do it yet.\n","metadata":{}}]}