diff --git a/README.md b/README.md index 261f617..1ae8091 100644 --- a/README.md +++ b/README.md @@ -64,9 +64,12 @@ Pretrain the exposure encoder as a denoising autoencoder using training-set EIDs ```bash python train_exposure_autoencoder.py \ --exposure_cache_dir ukb_exposure_cache \ - --train_eid_file ukb_train_eid.csv + --train_eid_file ukb_train_eid.csv \ + --val_eid_file ukb_val_eid.csv ``` The best checkpoint contains both `model_state_dict`, an `encoder_state_dict` compatible with the default gated `TimesNetExposureEncoder`, and the channel normalization statistics needed when the encoder is attached to DeepHealth. +Multi-GPU pretraining follows the main trainer interface: add +`--data_parallel --gpu_ids 0,1,2,3`. diff --git a/train_exposure_autoencoder.py b/train_exposure_autoencoder.py index 5f7418e..c3204fa 100644 --- a/train_exposure_autoencoder.py +++ b/train_exposure_autoencoder.py @@ -50,9 +50,9 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--exposure_cache_dir", required=True) parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") + parser.add_argument("--val_eid_file", default="ukb_val_eid.csv") parser.add_argument("--runs_root", default="runs") parser.add_argument("--seed", type=int, default=42) - parser.add_argument("--val_fraction", type=float, default=0.05) parser.add_argument("--n_embd", type=int, default=120) parser.add_argument("--d_model", type=int, default=None) parser.add_argument("--n_layers", type=int, default=2) @@ -72,35 +72,80 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--device", default="cuda") parser.add_argument("--amp", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument( + "--data_parallel", + action="store_true", + help="Use torch.nn.DataParallel across multiple CUDA devices.", + ) + parser.add_argument( + "--gpu_ids", + default=None, + help="Comma-separated CUDA device ids for --data_parallel, e.g. 0,1,2,3.", + ) args = parser.parse_args() - if not 0.0 < args.val_fraction < 1.0: - parser.error("--val_fraction must be between 0 and 1") if not 0.0 <= args.mask_ratio < 1.0: parser.error("--mask_ratio must be in [0, 1)") + if args.gpu_ids: + try: + args.gpu_ids = [ + int(part.strip()) + for part in args.gpu_ids.split(",") + if part.strip() + ] + except ValueError as exc: + parser.error("--gpu_ids must be a comma-separated list of integers") + if not args.gpu_ids: + parser.error("--gpu_ids did not contain any valid CUDA device ids") + args.data_parallel = True return args -def select_rows( - cache: ExposureCache, train_eids: set[int], val_fraction: float, seed: int -) -> tuple[np.ndarray, np.ndarray]: - candidate_eids = np.asarray( - sorted(set(map(int, cache.eids)) & train_eids), dtype=np.int64 - ) - if len(candidate_eids) < 2: - raise ValueError("Need at least two training EIDs with cached exposure") - rng = np.random.default_rng(seed) - rng.shuffle(candidate_eids) - n_val = max(1, int(round(len(candidate_eids) * val_fraction))) - val_eids = candidate_eids[:n_val] - fit_eids = candidate_eids[n_val:] +def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray: valid_row = np.asarray(cache.row_index, dtype=np.int64) >= 0 - fit_event_rows = valid_row & np.isin(cache.eids, fit_eids) - val_event_rows = valid_row & np.isin(cache.eids, val_eids) - fit_rows = np.unique(np.asarray(cache.row_index[fit_event_rows], dtype=np.int64)) - val_rows = np.unique(np.asarray(cache.row_index[val_event_rows], dtype=np.int64)) - if len(fit_rows) == 0 or len(val_rows) == 0: - raise ValueError("Training/validation exposure rows are empty after filtering") - return fit_rows, val_rows + selected_events = valid_row & np.isin(cache.eids, np.fromiter(eids, np.int64)) + rows = np.unique( + np.asarray(cache.row_index[selected_events], dtype=np.int64) + ) + if len(rows) == 0: + raise ValueError(f"{split} exposure rows are empty after EID filtering") + return rows + + +def maybe_wrap_data_parallel( + model: TimesNetExposureAutoencoder, + args: argparse.Namespace, + device: torch.device, + logger, +): + if not args.data_parallel: + return model + if device.type != "cuda": + raise ValueError("--data_parallel requires --device cuda or cuda:") + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + raise ValueError("--data_parallel requires at least two CUDA devices") + primary = ( + int(device.index) + if device.index is not None + else int(torch.cuda.current_device()) + ) + device_ids = ( + args.gpu_ids + if args.gpu_ids + else list(range(torch.cuda.device_count())) + ) + device_ids = [primary, *[idx for idx in device_ids if idx != primary]] + if len(device_ids) < 2: + raise ValueError("--data_parallel needs at least two device ids") + if any(idx < 0 or idx >= torch.cuda.device_count() for idx in device_ids): + raise ValueError(f"CUDA device id is out of range: {device_ids}") + logger.info(f"Using DataParallel on CUDA devices: {device_ids}") + return torch.nn.DataParallel( + model, device_ids=device_ids, output_device=primary + ) + + +def unwrap_model(model) -> TimesNetExposureAutoencoder: + return model.module if isinstance(model, torch.nn.DataParallel) else model def channel_stats( @@ -133,7 +178,7 @@ def masked_mse( def run_epoch( - model: TimesNetExposureAutoencoder, + model, loader: DataLoader, device: torch.device, stats: tuple[torch.Tensor, ...], @@ -218,9 +263,12 @@ def main() -> None: ) logger = setup_logging(run_dir) cache = ExposureCache(args.exposure_cache_dir) - train_rows, val_rows = select_rows( - cache, load_eid_file(args.train_eid_file), args.val_fraction, args.seed - ) + train_eids = load_eid_file(args.train_eid_file) + val_eids = load_eid_file(args.val_eid_file) + if train_eids & val_eids: + raise ValueError("train and validation EID files must be disjoint") + train_rows = select_rows(cache, train_eids, "Training") + val_rows = select_rows(cache, val_eids, "Validation") raw_stats = channel_stats(cache, train_rows) stats = tuple( torch.as_tensor(value, device=device).view(1, 1, -1) @@ -244,6 +292,7 @@ def main() -> None: conv_kernel_size=args.conv_kernel_size, mlp_ratio=args.mlp_ratio, dropout=args.dropout, ).to(device) + model = maybe_wrap_data_parallel(model, args, device, logger) optimizer = AdamW( model.parameters(), lr=args.base_lr, weight_decay=args.weight_decay, betas=(0.9, 0.95), @@ -292,10 +341,11 @@ def main() -> None: if val_loss < best_loss: best_loss = val_loss stale_epochs = 0 + checkpoint_model = unwrap_model(model) torch.save( { - "model_state_dict": model.state_dict(), - "encoder_state_dict": model.encoder.state_dict(), + "model_state_dict": checkpoint_model.state_dict(), + "encoder_state_dict": checkpoint_model.encoder.state_dict(), "model_config": { key: config[key] for key in ( "n_embd", "d_model", "n_layers", "top_k",