JiaruiLi/DeepHealth
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Avoid multiprocessing data loader for shape export

Browse Source
This commit is contained in:
Jiarui Li
2026-07-01 15:37:36 +08:00
parent d08e5b34f4
commit bf31aa0432
2 changed files with 29 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
export_weibull_death_parameter_stats.py
View File

@@ -18,6 +18,7 @@ import numpy as np
import pandas as pd import pandas as pd
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import torch.multiprocessing as torch_mp
from torch.utils.data import DataLoader from torch.utils.data import DataLoader
from tqdm.auto import tqdm from tqdm.auto import tqdm
@@ -40,6 +41,11 @@ from evaluate_auc_v2 import (
validate_dataset_metadata, 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: 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] 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("--landmark_step", type=float, default=None)
parser.add_argument("--horizons", type=str, default=None) parser.add_argument("--horizons", type=str, default=None)
parser.add_argument("--batch_size", type=int, 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("--device", type=str, default=None)
parser.add_argument("--use_amp", action=argparse.BooleanOptionalAction, 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"]) 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)) batch_size = int(cfg_get(args, cfg, "batch_size", 128))
num_workers = int(cfg_get(args, cfg, "num_workers", 4)) num_workers = int(cfg_get(args, cfg, "num_workers", 0))
loader = DataLoader( loader_kwargs = {
landmark_dataset, "batch_size": batch_size,
batch_size=batch_size, "shuffle": False,
shuffle=False, "collate_fn": collate_landmark_fn,
collate_fn=collate_landmark_fn, "num_workers": num_workers,
num_workers=num_workers, "pin_memory": device.type == "cuda",
pin_memory=device.type == "cuda", }
persistent_workers=num_workers > 0, if num_workers > 0:
prefetch_factor=2 if num_workers > 0 else None, 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)) use_amp = bool(cfg_get(args, cfg, "use_amp", False))
hidden_all, row_arrays = infer_landmark_hidden_local( hidden_all, row_arrays = infer_landmark_hidden_local(

4
run_weibull_shape_exports.sh
View File

@@ -9,7 +9,7 @@
# bash run_weibull_shape_exports.sh # bash run_weibull_shape_exports.sh
# #
# Optional overrides: # 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 # bash run_weibull_shape_exports.sh
set -euo pipefail set -euo pipefail
@@ -17,7 +17,7 @@ set -euo pipefail
PYTHON="${PYTHON:-python}" PYTHON="${PYTHON:-python}"
DEVICE="${DEVICE:-cuda}" DEVICE="${DEVICE:-cuda}"
BATCH_SIZE="${BATCH_SIZE:-128}" BATCH_SIZE="${BATCH_SIZE:-128}"
NUM_WORKERS="${NUM_WORKERS:-4}" NUM_WORKERS="${NUM_WORKERS:-0}"
ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}" ROW_BATCH_SIZE="${ROW_BATCH_SIZE:-512}"
LANDMARK_START="${LANDMARK_START:-40}" LANDMARK_START="${LANDMARK_START:-40}"
LANDMARK_STOP="${LANDMARK_STOP:-80}" LANDMARK_STOP="${LANDMARK_STOP:-80}"