diff --git a/export_weibull_death_parameter_stats.py b/export_weibull_death_parameter_stats.py index 79914cf..dd1eef2 100644 --- a/export_weibull_death_parameter_stats.py +++ b/export_weibull_death_parameter_stats.py @@ -18,6 +18,7 @@ import numpy as np import pandas as pd import torch import torch.nn.functional as F +import torch.multiprocessing as torch_mp from torch.utils.data import DataLoader from tqdm.auto import tqdm @@ -40,6 +41,11 @@ from evaluate_auc_v2 import ( validate_dataset_metadata, ) +try: + torch_mp.set_sharing_strategy("file_system") +except RuntimeError: + pass + def quantile_summary(df: pd.DataFrame, group_cols: List[str], value_cols: List[str]) -> pd.DataFrame: probs = [0.01, 0.05, 0.25, 0.50, 0.75, 0.95, 0.99] @@ -293,7 +299,15 @@ def main() -> None: parser.add_argument("--landmark_step", type=float, default=None) parser.add_argument("--horizons", type=str, default=None) parser.add_argument("--batch_size", type=int, default=None) - parser.add_argument("--num_workers", type=int, default=None) + parser.add_argument( + "--num_workers", + type=int, + default=0, + help=( + "DataLoader workers. Default 0 avoids Linux multiprocessing " + "'received 0 items of ancdata' failures on shared filesystems." + ), + ) parser.add_argument("--device", type=str, default=None) parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, default=None) parser.add_argument("--hidden_cache_dtype", type=str, default="float32", choices=["float16", "float32"]) @@ -386,17 +400,18 @@ def main() -> None: ) batch_size = int(cfg_get(args, cfg, "batch_size", 128)) - num_workers = int(cfg_get(args, cfg, "num_workers", 4)) - loader = DataLoader( - landmark_dataset, - batch_size=batch_size, - shuffle=False, - collate_fn=collate_landmark_fn, - num_workers=num_workers, - pin_memory=device.type == "cuda", - persistent_workers=num_workers > 0, - prefetch_factor=2 if num_workers > 0 else None, - ) + num_workers = int(cfg_get(args, cfg, "num_workers", 0)) + loader_kwargs = { + "batch_size": batch_size, + "shuffle": False, + "collate_fn": collate_landmark_fn, + "num_workers": num_workers, + "pin_memory": device.type == "cuda", + } + if num_workers > 0: + loader_kwargs["persistent_workers"] = True + loader_kwargs["prefetch_factor"] = 2 + loader = DataLoader(landmark_dataset, **loader_kwargs) use_amp = bool(cfg_get(args, cfg, "use_amp", False)) hidden_all, row_arrays = infer_landmark_hidden_local( diff --git a/run_weibull_shape_exports.sh b/run_weibull_shape_exports.sh index 3bd46e5..eec5b82 100755 --- a/run_weibull_shape_exports.sh +++ b/run_weibull_shape_exports.sh @@ -9,7 +9,7 @@ # bash run_weibull_shape_exports.sh # # Optional overrides: -# PYTHON=python3 DEVICE=cuda BATCH_SIZE=128 NUM_WORKERS=4 ROW_BATCH_SIZE=512 \ +# PYTHON=python3 DEVICE=cuda BATCH_SIZE=128 NUM_WORKERS=0 ROW_BATCH_SIZE=512 \ # bash run_weibull_shape_exports.sh set -euo pipefail @@ -17,7 +17,7 @@ set -euo pipefail PYTHON="${PYTHON:-python}" DEVICE="${DEVICE:-cuda}" BATCH_SIZE="${BATCH_SIZE:-128}" -NUM_WORKERS="${NUM_WORKERS:-4}" +NUM_WORKERS="${NUM_WORKERS:-0}" ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}" LANDMARK_START="${LANDMARK_START:-40}" LANDMARK_STOP="${LANDMARK_STOP:-80}"