OSCR

Robust disease prognosis via diagnostic knowledge preservation: A sequential learning approach.

Code ↔ Paper

5 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 5 matches
  1. [1] § 2. Materials and methods › 2.5. Experimental setup › 2.5.2. Training protocol. ↔ Breast_Cancer/Multitask/train.py, lines 592–660 · score 0.77 · cross entropy loss, learning rate scheduler, BI RADS, Adam, breast cancer, optimizer
  2. [2] § 2. Materials and methods › 2.5. Experimental setup › 2.5.3. Evaluation and analysis. ↔ Breast_Cancer/Single_Task/eval_birads.py, lines 108–197 · score 0.76 · Micro AUROC, Macro AUROC, multi class, Balanced accuracy, predictions
  3. [3] § 2. Materials and methods › 2.5. Experimental setup › 2.5.3. Evaluation and analysis. ↔ Breast_Cancer/Multitask/evaluate.py, lines 28–113 · score 0.70 · Macro AUROC, multi class, Balanced accuracy, sensitivity, micro, predictions
  4. [4] § 2. Materials and methods › 2.5. Experimental setup › 2.5.2. Training protocol. ↔ Knee_OA/train.py, lines 90–177 · score 0.67 · cosine annealing learning, Adam, scheduler, augmentation, optimizer, loss
  5. [5] § 2. Materials and methods › 2.5. Experimental setup › 2.5.2. Training protocol. ↔ Breast_Cancer/Multitask/train.py, lines 592–660 · score 0.52 · prognosis AUROC, multitask learning, epoch, probability, loss, replay

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 796 lines · 31 KB · CC-BY-NC-4.0 · 2 matches

  1. import os
  2. import wandb
  3. import torch
  4. import torch.nn as nn
  5. import torch.optim as optim
  6. import torch.nn.functional as F
  7. from torch.utils.data import DataLoader, Dataset
  8. from sklearn.metrics import roc_auc_score, average_precision_score
  9. import numpy as np
  10. from tqdm import tqdm
  11. import argparse
  12. import pickle
  14. from models_torch import BaselineBreastModel
  15. from torch_transforms import compose_transform
  16. import load_mammogram
  18. os.environ['WANDB_API_KEY'] = #
  20. class MammogramDataset(Dataset):
  21. """Dataset class for mammogram images with multi-task support"""
  22. def __init__(self, img_dir, datalist, transform, crop_size=(2944, 1920), tasks=None):
  23. self.img_dir = img_dir
  24. self.datalist = datalist
  25. self.transform = transform
  26. self.crop_size = crop_size
  27. self.tasks = tasks if tasks else ["prognosis"] # Default to single task if not specified
  29. # Filter out samples with missing labels for the required tasks
  30. self.valid_indices = []
  31. for idx, metadata in enumerate(self.datalist):
  32. is_valid = True
  34. # Check if BI-RADS label is valid if birads task is requested
  35. if "birads" in self.tasks:
  36. birads_value = metadata.get('birads', None)
  37. if birads_value is None or birads_value == 'NA' or birads_value == 'inconclusive':
  38. is_valid = False
  40. if is_valid:
  41. self.valid_indices.append(idx)
  43. print(f"Dataset initialized with {len(self.valid_indices)}/{len(datalist)} valid samples for tasks: {tasks}")
  45. def __len__(self):
  46. return len(self.valid_indices)
  48. def __getitem__(self, idx):
  49. original_idx = self.valid_indices[idx]
  50. metadata = self.datalist[original_idx]
  52. # Load mammogram exam
  53. images, scans_info = load_mammogram.load_mammogram_exam(
  54. img_dir=self.img_dir,
  55. metadata=metadata,
  56. transform=self.transform,
  57. select_index_logic=("random", 4),
  58. crop_method="best_center",
  59. crop_size=self.crop_size,
  60. skip_img=False
  61. )
  63. # Prepare image data
  64. x = {}
  65. for i, scan in enumerate(scans_info):
  66. side = scan['laterality']
  67. if side == "left":
  68. x[f"L-{scan['view']}"] = images[i,0:1,:,:]
  69. else:
  70. x[f"R-{scan['view']}"] = images[i,0:1,:,:]
  72. # Prepare labels for all requested tasks
  73. labels = {}
  74. if "prognosis" in self.tasks:
  75. labels["prognosis"] = metadata['5yr_cancer_label']['malignant']
  77. if "birads" in self.tasks and 'birads' in metadata and metadata['birads'] not in [None, 3, 4, 5]:
  78. labels["birads"] = metadata['birads']
  79. else:
  80. labels["birads"] = -1
  82. # Return different outputs based on the number of tasks
  83. if len(self.tasks) == 1:
  84. # For single task, return a single label for backward compatibility
  85. return x, labels[self.tasks[0]]
  86. else:
  87. # For multi-task, return a dictionary of labels
  88. return x, labels
  90. def collate_mammogram_batch(batch):
  91. """
  92. Custom collate function for mammogram batches with single or multiple labels per sample.
  94. Args:
  95. batch: List of tuples (x, label) where:
  96. x: Dictionary with keys 'L-CC', 'L-MLO', 'R-CC', 'R-MLO' mapping to image tensors
  97. label: Either a single value or a dictionary of task labels
  99. Returns:
  100. Tuple of:
  101. batch_x: Dictionary mapping view names to batched image tensors
  102. batch_labels: Dictionary of label tensors or single tensor depending on input format
  103. """
  104. # Initialize dictionary for images
  105. batch_x = {
  106. 'L-CC': [], 'L-MLO': [],
  107. 'R-CC': [], 'R-MLO': []
  108. }
  110. # Check if using multi-task labels (dictionary) or single task label
  111. first_label = batch[0][1]
  112. is_multi_task = isinstance(first_label, dict)
  114. if is_multi_task:
  115. # Initialize dictionary for each task
  116. tasks = first_label.keys()
  117. batch_labels = {task: [] for task in tasks}
  119. # Collect images and task labels
  120. for x, labels in batch:
  121. for view in batch_x.keys():
  122. if view in x:
  123. batch_x[view].append(x[view])
  124. else:
  125. # Handle missing views with a tensor of zeros
  126. # Assuming all tensors in a batch have the same shape
  127. shape = next(iter(x.values())).shape
  128. batch_x[view].append(torch.zeros(shape))
  130. for task, label in labels.items():
  131. batch_labels[task].append(label)
  133. # Stack tensors for each view and create label tensors
  134. batch_x = {view: torch.stack(tensors) for view, tensors in batch_x.items()}
  135. batch_labels = {task: torch.tensor(labels) for task, labels in batch_labels.items()}
  136. else:
  137. # Single task mode - for backward compatibility
  138. batch_labels = []
  140. # Collect images and labels
  141. for x, label in batch:
  142. for view in batch_x.keys():
  143. if view in x:
  144. batch_x[view].append(x[view])
  145. else:
  146. # Handle missing views with a tensor of zeros
  147. shape = next(iter(x.values())).shape
  148. batch_x[view].append(torch.zeros(shape))
  149. batch_labels.append(label)
  151. # Stack tensors for each view and create label tensor
  152. batch_x = {view: torch.stack(tensors) for view, tensors in batch_x.items()}
  153. batch_labels = torch.tensor(batch_labels)
  155. return batch_x, batch_labels
  157. def unpickle_from_file(file_name):
  158. with open(file_name, 'rb') as handle:
  159. try:
  160. return pickle.load(handle)
  161. except ImportError:
  162. return pd.read_pickle(file_name)
  164. class MultiTaskBreastModel(nn.Module):
  165. """Wrapper for baseline breast model to support multiple output heads"""
  166. def __init__(self, base_model, device, tasks=None):
  167. super(MultiTaskBreastModel, self).__init__()
  168. self.tasks = tasks if tasks else ["prognosis"]
  170. self.base_model = base_model
  171. original_fc2 = self.base_model.fc2
  175. # Base model
  178. # Get feature dimension from base model
  179. self.feature_dim = self.base_model.fc2.in_features
  181. # Create task-specific output heads
  182. self.heads = nn.ModuleDict()
  184. if "prognosis" in self.tasks:
  185. self.heads["prognosis"] = nn.Linear(self.feature_dim, 1)
  187. if "birads" in self.tasks:
  188. self.heads["birads"] = nn.Linear(original_fc2.in_features, original_fc2.out_features)
  189. # Optionally copy weights and biases
  190. self.heads["birads"].weight.data = original_fc2.weight.data.clone()
  191. self.heads["birads"].bias.data = original_fc2.bias.data.clone()
  193. self.base_model.fc2 = nn.Identity()
  195. def forward(self, x):
  196. # Extract features from base model
  197. features = self.base_model(x)
  199. # Apply task-specific heads
  200. outputs = {}
  201. for task in self.tasks:
  202. outputs[task] = self.heads[task](features)
  204. # Reshape binary outputs to match expected format
  205. if task == "prognosis":
  206. outputs[task] = outputs[task].squeeze(-1)
  208. return outputs
  211. def get_non_requires_grad_params(model):
  212. """
  213. Find all parameters in a model that do NOT require gradients.
  215. Args:
  216. model: PyTorch model
  218. Returns:
  219. List of parameter names that don't require gradients and their shapes
  220. """
  221. non_requires_grad_params = []
  223. for name, param in model.named_parameters():
  224. if not param.requires_grad:
  225. non_requires_grad_params.append((name, param.shape))
  227. return non_requires_grad_params
  230. def load_pretrained_weights(model, pretrained_path):
  231. """Load pretrained BI-RADS weights and modify for binary classification"""
  232. if pretrained_path:
  233. pretrained_dict = torch.load(pretrained_path, map_location='cpu')
  234. model_dict = model.state_dict()
  236. # Filter out fc2 layer from pretrained weights
  237. pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict and 'fc2' not in k}
  239. # Update model weights
  240. model_dict.update(pretrained_dict)
  241. model.load_state_dict(model_dict)
  242. return model
  244. def train_epoch(model, train_loader, criterions, optimizer, device, tasks, diagnosis_loader=None, replay_prob=0.5):
  245. print(tasks)
  246. model.train()
  247. running_loss = 0.0
  248. task_running_losses = {task: 0.0 for task in tasks}
  250. # For metrics tracking - only track prognosis for simplicity
  251. task_metrics = {}
  252. if 'prognosis' in tasks:
  253. task_metrics['prognosis'] = {
  254. 'predictions': [],
  255. 'targets': []
  256. }
  258. if diagnosis_loader is not None:
  259. task_running_losses['birads'] = 0.0
  261. # Set up diagnosis iterator if provided
  262. diagnosis_iterator = iter(diagnosis_loader) if diagnosis_loader else None
  264. pbar = tqdm(train_loader, desc='Training')
  265. for batch_idx, (data, targets) in enumerate(pbar):
  266. # Determine whether to do replay for this batch
  267. do_replay = diagnosis_loader is not None and np.random.uniform() <= replay_prob
  269. if do_replay:
  270. try:
  271. # Get a batch from diagnosis dataset
  272. print("its here")
  273. replay_data, replay_targets = next(diagnosis_iterator)
  275. # Replace current batch with diagnosis data
  276. data = replay_data
  278. # Handle targets format
  279. if isinstance(replay_targets, dict):
  280. targets = replay_targets
  281. else:
  282. # Convert to dict format if needed
  283. targets = {"birads": replay_targets}
  285. # Set current tasks for this batch
  286. current_tasks = ["birads"]
  288. except StopIteration:
  289. # Reset iterator if we've gone through all diagnosis data
  290. diagnosis_iterator = iter(diagnosis_loader)
  291. replay_data, replay_targets = next(diagnosis_iterator)
  293. # Replace current batch with diagnosis data
  294. data = replay_data
  296. # Handle targets format
  297. if isinstance(replay_targets, dict):
  298. targets = replay_targets
  299. else:
  300. # Convert to dict format if needed
  301. targets = {"birads": replay_targets}
  303. # Set current tasks for this batch
  304. current_tasks = ["birads"]
  305. else:
  306. # Regular prognosis batch
  307. current_tasks = tasks
  309. # Move data to device
  310. for view in data:
  311. data[view] = data[view].to(device)
  313. # Handle targets based on format
  314. if not isinstance(targets, dict):
  315. # Convert single task target to dict format for consistent handling
  316. # Assuming single task is always the first task in the list
  317. targets = {current_tasks[0]: targets.to(device)}
  318. else:
  319. # Move all targets to device
  320. for task in targets:
  321. if task in targets: # Only process tasks that have labels
  322. targets[task] = targets[task].to(device)
  324. # Convert birads to long for CrossEntropyLoss
  325. if task == "birads":
  326. targets[task] = targets[task].long()
  328. optimizer.zero_grad()
  329. outputs = model(data)
  331. # Calculate task-specific losses and total loss
  332. total_loss = 0.0
  333. for task in current_tasks:
  334. if task in targets and task in outputs:
  335. mask = (targets[task] != -1)
  337. if isinstance(mask, bool):
  338. if not mask:
  339. continue
  340. mask = torch.tensor([True]) # Create a single-element boolean tensor
  342. elif mask.sum() == 0:
  343. continue
  345. if isinstance(criterions[task], nn.BCEWithLogitsLoss):
  346. # For binary cross entropy, both outputs and targets should be float
  347. task_loss = criterions[task](
  348. outputs[task][mask].float(),
  349. targets[task][mask].float()
  350. )
  351. elif isinstance(criterions[task], nn.CrossEntropyLoss):
  352. # For cross entropy, outputs should be float but targets should be long
  353. task_loss = criterions[task](
  354. outputs[task][mask].float(),
  355. targets[task][mask].long()
  356. )
  357. # print("birads targets ", targets[task][mask].long())
  358. else:
  359. # For other loss functions, maintain original types
  360. task_loss = criterions[task](outputs[task][mask], targets[task][mask])
  363. task_running_losses[task] += task_loss.item()
  364. total_loss += task_loss
  366. # Store predictions for prognosis metrics only (if it's a prognosis batch)
  367. if task == 'prognosis' and task in task_metrics and task in current_tasks:
  368. task_metrics[task]['predictions'].extend(torch.sigmoid(outputs[task]).detach().cpu().numpy())
  369. task_metrics[task]['targets'].extend(targets[task].cpu().numpy())
  371. total_loss.backward()
  372. optimizer.step()
  374. running_loss += total_loss.item()
  376. # Update progress bar
  377. pbar.set_postfix({
  378. 'loss': total_loss.item(),
  379. **{f"{task}_loss": task_running_losses[task]/(batch_idx+1) for task in tasks if task in task_running_losses}
  380. })
  382. # Calculate average losses
  383. epoch_loss = running_loss / len(train_loader)
  384. task_losses = {task: task_running_losses[task] / len(train_loader) for task in tasks if task in task_running_losses}
  386. # Calculate metrics for prognosis task only
  387. metrics = {}
  388. if 'prognosis' in task_metrics and len(task_metrics['prognosis']['predictions']) > 0:
  389. preds = np.array(task_metrics['prognosis']['predictions'])
  390. targets = np.array(task_metrics['prognosis']['targets'])
  392. metrics["prognosis_auroc"] = roc_auc_score(targets, preds)
  393. metrics["prognosis_auprc"] = average_precision_score(targets, preds)
  395. return epoch_loss, task_losses, metrics
  397. def validate(model, val_loader, criterions, device, tasks):
  398. model.eval()
  399. running_loss = 0.0
  400. task_running_losses = {task: 0.0 for task in tasks}
  402. # For metrics tracking - only track prognosis
  403. task_metrics = {}
  404. if 'prognosis' in tasks:
  405. task_metrics['prognosis'] = {
  406. 'predictions': [],
  407. 'targets': []
  408. }
  410. with torch.no_grad():
  411. for data, targets in tqdm(val_loader, desc="Validation"):
  412. # Move data to device
  413. for view in data:
  414. data[view] = data[view].to(device)
  416. # Handle both single-task and multi-task targets
  417. if not isinstance(targets, dict):
  418. # Convert single task target to dict format for consistent handling
  419. targets = {"prognosis": targets.to(device)}
  420. else:
  421. # Move all targets to device
  422. for task in targets:
  423. if task in targets: # Only process tasks that have labels
  424. targets[task] = targets[task].to(device)
  426. # Convert birads to long for CrossEntropyLoss
  427. if task == "birads":
  428. targets[task] = targets[task].long()
  430. outputs = model(data)
  432. # Calculate task-specific losses and total loss
  433. total_loss = 0.0
  434. for task in tasks:
  435. if task in targets and task in outputs:
  436. mask = (targets[task] != -1)
  438. if isinstance(mask, bool):
  439. if not mask:
  440. continue
  441. mask = torch.tensor([True]) # Create a single-element boolean tensor
  443. elif mask.sum() == 0:
  444. continue
  446. if isinstance(criterions[task], nn.BCEWithLogitsLoss):
  447. # For binary cross entropy, both outputs and targets should be float
  448. task_loss = criterions[task](
  449. outputs[task][mask].float(),
  450. targets[task][mask].float()
  451. )
  452. elif isinstance(criterions[task], nn.CrossEntropyLoss):
  453. # For cross entropy, outputs should be float but targets should be long
  454. task_loss = criterions[task](
  455. outputs[task][mask].float(),
  456. targets[task][mask].long()
  457. )
  458. else:
  459. # For other loss functions, maintain original types
  460. task_loss = criterions[task](outputs[task][mask], targets[task][mask])
  462. task_running_losses[task] += task_loss.item()
  463. total_loss += task_loss
  465. # Store predictions for prognosis metrics only
  466. if task == 'prognosis' and task in task_metrics:
  467. task_metrics[task]['predictions'].extend(torch.sigmoid(outputs[task]).detach().cpu().numpy())
  468. task_metrics[task]['targets'].extend(targets[task].cpu().numpy())
  470. running_loss += total_loss.item()
  472. # Calculate average losses
  473. epoch_loss = running_loss / len(val_loader)
  474. task_losses = {task: task_running_losses[task] / len(val_loader) for task in tasks if task in task_running_losses}
  476. # Calculate metrics for prognosis task only
  477. metrics = {}
  478. if 'prognosis' in task_metrics and len(task_metrics['prognosis']['predictions']) > 0:
  479. preds = np.array(task_metrics['prognosis']['predictions'])
  480. targets = np.array(task_metrics['prognosis']['targets'])
  482. metrics["prognosis_auroc"] = roc_auc_score(targets, preds)
  483. metrics["prognosis_auprc"] = average_precision_score(targets, preds)
  485. return epoch_loss, task_losses, metrics
  487. def train(args):
  488. wandb.init(project=args.wandb_project, name=args.wandb_run_name)
  489. wandb.config.update(args)
  491. device = torch.device("cuda" if args.device_type == "gpu" else "cpu")
  493. # Define tasks based on arguments
  494. tasks = ["prognosis"]
  495. if args.multitask:
  496. tasks.append("birads")
  498. print(f"Training with tasks: {tasks}")
  500. # Load prognosis data
  501. prognosis_train_data = unpickle_from_file(f"{args.prognosis_datalist_path}/Fold_{args.fold}_train.pkl")
  502. val_data = unpickle_from_file(f"{args.prognosis_datalist_path}/Fold_{args.fold}_val.pkl")
  503. # np.random.shuffle(prognosis_train_data)
  504. # prognosis_train_data = prognosis_train_data[:256]
  505. # np.random.shuffle(val_data)
  506. # val_data = val_data[:256]
  507. transform = compose_transform(augmentation="standard", resize=None, image_format="greyscale")
  509. # Create primary dataset for prognosis
  510. train_dataset = MammogramDataset(
  511. img_dir=args.img_dir,
  512. datalist=prognosis_train_data,
  513. transform=transform,
  514. crop_size=args.crop_size,
  515. tasks=tasks # Primary dataset includes all primary tasks
  516. )
  518. val_dataset = MammogramDataset(
  519. img_dir=args.img_dir,
  520. datalist=val_data,
  521. transform=compose_transform(image_format="greyscale"),
  522. crop_size=args.crop_size,
  523. tasks=["prognosis", "birads"] # Validation dataset includes all tasks
  524. )
  526. # Create data loaders for prognosis
  527. train_loader = DataLoader(
  528. train_dataset,
  529. batch_size=args.batch_size,
  530. shuffle=True,
  531. num_workers=args.num_workers,
  532. collate_fn=collate_mammogram_batch,
  533. pin_memory=True
  534. )
  536. val_loader = DataLoader(
  537. val_dataset,
  538. batch_size=args.batch_size,
  539. shuffle=False,
  540. num_workers=args.num_workers,
  541. collate_fn=collate_mammogram_batch,
  542. pin_memory=True
  543. )
  545. # Create diagnosis dataset (separate cohort)
  546. diagnosis_loader = None
  547. if args.replay:
  548. print("Setting up diagnosis loader for batch-level replay")
  549. # Load diagnosis data from a separate path
  550. diagnosis_train_data = unpickle_from_file(f"{args.diagnosis_datalist_path}/train.pkl")
  552. diagnosis_dataset = MammogramDataset(
  553. img_dir=args.img_dir,
  554. datalist=diagnosis_train_data,
  555. transform=transform,
  556. crop_size=args.crop_size,
  557. tasks=["birads"] # This dataset is only for diagnosis
  558. )
  560. diagnosis_loader = DataLoader(
  561. diagnosis_dataset,
  562. batch_size=args.batch_size,
  563. shuffle=True,
  564. num_workers=args.num_workers,
  565. collate_fn=collate_mammogram_batch,
  566. pin_memory=True
  567. )
  569. nodropout_probability = 1 - args.dropout_p
  570. base_model = BaselineBreastModel(
  571. device=device,
  572. nodropout_probability=nodropout_probability,
  573. gaussian_noise_std=args.gaussian_noise_std,
  574. )
  576. if args.pretrained_path:
  577. base_model = load_pretrained_weights(base_model, args.pretrained_path)
  579. model = MultiTaskBreastModel(
  580. base_model=base_model,
  581. device=device,
  582. tasks=["prognosis", "birads"] # Always include both tasks in model
  583. ).to(device)
  585. frozen_params = get_non_requires_grad_params(model)
  588. print(f"Total parameters NOT requiring gradients: {len(frozen_params)}")
  589. for name, shape in frozen_params:
  590. print(f"{name}: {shape}")
  592. # Set up class weights for BI-RADS using the provided counts
  593. class_counts = torch.tensor([6634, 17388,15939], dtype=torch.float).to(device)
  594. class_weights = 1.0 / class_counts
  595. class_weights = class_weights / class_weights.sum() * len(class_counts)
  597. # Setup loss functions for each task
  598. criterions = {
  599. "prognosis": nn.BCEWithLogitsLoss(),
  600. "birads": nn.CrossEntropyLoss(weight=class_weights)
  601. }
  603. optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
  605. scheduler = optim.lr_scheduler.ReduceLROnPlateau(
  606. optimizer, mode='max', factor=0.75, patience=10, verbose=True
  607. )
  609. # Create save directory
  610. save_path = os.path.join(args.save_path, f"Fold_{args.fold}", args.wandb_run_name)
  611. os.makedirs(save_path, exist_ok=True)
  613. # Track best metrics for each task
  614. best_metrics = {
  615. "total_loss": float('inf'),
  616. "prognosis_auroc": 0,
  617. "prognosis_auprc": 0,
  618. }
  620. patience_counter = 0
  622. # Training loop
  623. for epoch in range(args.epochs):
  624. print(f'\nEpoch {epoch+1}/{args.epochs}')
  626. # Training phase - with batch-level replay
  627. if args.replay and diagnosis_loader is not None:
  628. # Use new train_epoch function with batch-level replay
  629. total_train_loss, all_train_task_losses, all_train_metrics = train_epoch(
  630. model=model,
  631. train_loader=train_loader,
  632. criterions=criterions,
  633. optimizer=optimizer,
  634. device=device,
  635. tasks=["prognosis"],
  636. diagnosis_loader=diagnosis_loader,
  637. replay_prob=args.replay_prob
  638. )
  639. # Joint multitask learning
  640. elif args.multitask:
  641. # Train on both prognosis and birads tasks in the same dataset
  642. total_train_loss, all_train_task_losses, all_train_metrics = train_epoch(
  643. model=model,
  644. train_loader=train_loader,
  645. criterions=criterions,
  646. optimizer=optimizer,
  647. device=device,
  648. tasks=["prognosis", "birads"]
  649. )
  650. # Single task mode (prognosis only)
  651. else:
  652. # Standard single-task training
  653. total_train_loss, all_train_task_losses, all_train_metrics = train_epoch(
  654. model=model,
  655. train_loader=train_loader,
  656. criterions=criterions,
  657. optimizer=optimizer,
  658. device=device,
  659. tasks=["prognosis"]
  660. )
  662. # Validation phase - always validate on all tasks
  663. val_loss, val_task_losses, val_metrics = validate(
  664. model=model,
  665. val_loader=val_loader,
  666. criterions=criterions,
  667. device=device,
  668. tasks=["prognosis", "birads"]
  669. )
  671. # Log metrics
  672. log_dict = {
  673. 'epoch': epoch,
  674. 'train_loss': total_train_loss,
  675. 'val_loss': val_loss,
  676. **{f"train_{task}_loss": loss for task, loss in all_train_task_losses.items()},
  677. **{f"val_{task}_loss": loss for task, loss in val_task_losses.items()},
  678. **{f"train_{metric}": value for metric, value in all_train_metrics.items()},
  679. **{f"val_{metric}": value for metric, value in val_metrics.items()},
  680. 'learning_rate': optimizer.param_groups[0]['lr']
  681. }
  682. wandb.log(log_dict)
  684. # Print metrics
  685. print(f"Epoch {epoch+1} Results:")
  686. for metric, value in val_metrics.items():
  687. print(f" Val {metric}: {value:.4f}")
  689. # Update learning rate based on primary task metric (prognosis AUROC)
  690. if 'prognosis_auroc' in val_metrics:
  691. scheduler.step(val_metrics['prognosis_auroc'])
  693. # Check if this is the best model so far
  694. improved = False
  696. # Check each metric
  697. for metric in best_metrics:
  698. if metric == 'total_loss':
  699. if val_loss < best_metrics[metric]:
  700. best_metrics[metric] = val_loss
  701. improved = True
  702. torch.save(model.state_dict(), os.path.join(save_path, f'best_loss_model.pth'))
  703. elif metric in val_metrics and best_metrics[metric] is not None:
  704. if val_metrics[metric] > best_metrics[metric]:
  705. best_metrics[metric] = val_metrics[metric]
  706. improved = True
  707. torch.save(model.state_dict(), os.path.join(save_path, f'best_{metric}_model.pth'))
  709. # Also save a checkpoint for this epoch
  710. torch.save({
  711. 'epoch': epoch,
  712. 'model_state_dict': model.state_dict(),
  713. 'optimizer_state_dict': optimizer.state_dict(),
  714. 'scheduler_state_dict': scheduler.state_dict(),
  715. 'best_metrics': best_metrics,
  716. **val_metrics
  717. }, os.path.join(save_path, f'checkpoint_epoch_{epoch}.pth'))
  719. # Check for early stopping
  720. if improved:
  721. patience_counter = 0
  722. else:
  723. patience_counter += 1
  725. if patience_counter >= args.patience:
  726. print(f'Early stopping triggered after epoch {epoch+1}')
  727. break
  729. wandb.finish()
  731. # Save final model
  732. torch.save(model.state_dict(), os.path.join(save_path, 'final_model.pth'))
  734. # Print best results
  735. print("\nTraining completed. Best results:")
  736. for metric, value in best_metrics.items():
  737. if value is not None:
  738. print(f" Best {metric}: {value:.4f}")
  740. if __name__ == "__main__":
  741. parser = argparse.ArgumentParser(description='Train multi-task mammogram classification model')
  743. # Model parameters
  744. parser.add_argument('--pretrained-path', type=str, default=None,
  745. help='Path to pretrained model weights')
  746. parser.add_argument('--save-path', type=str, default=None,
  747. help='Path to save models')
  748. parser.add_argument('--gaussian-noise-std', type=float, default=0.01,
  749. help='Standard deviation of Gaussian noise')
  750. parser.add_argument('--dropout-p', type=float, default=0.1,
  751. help='Dropout probability')
  753. # Multi-task parameters
  754. parser.add_argument('--multitask', action='store_true',
  755. help='Enable multi-task learning with prognosis and diagnosis')
  756. parser.add_argument('--replay', action='store_true',
  757. help='Use experience replay for multi-task learning')
  758. parser.add_argument('--replay-prob', type=float, default=0.5,
  759. help='Probability of doing experience replay in each epoch')
  761. # Training parameters
  762. parser.add_argument('--batch-size', type=int, default=8)
  763. parser.add_argument('--learning-rate', type=float, default=1e-5)
  764. parser.add_argument('--epochs', type=int, default=100)
  765. parser.add_argument('--patience', type=int, default=10,
  766. help='Early stopping patience')
  767. parser.add_argument('--num-workers', type=int, default=8)
  769. # Device parameters
  770. parser.add_argument('--device-type', type=str, default='gpu')
  771. parser.add_argument('--gpu-number', type=int, default=0)
  773. # Prognosis data parameters
  774. parser.add_argument('--img-dir', type=str, required=True,
  775. help='Root directory containing the prognosis image data')
  776. parser.add_argument('--prognosis-datalist-path', type=str, required=True,
  777. help='Path to the pickle file containing prognosis train/val/test splits')
  779. # Diagnosis data parameters (for multitask or replay)
  780. parser.add_argument('--diagnosis-datalist-path', type=str,
  781. help='Path to the pickle file containing diagnosis train/val/test splits')
  783. # Common data parameters
  784. parser.add_argument('--crop-size', type=int, nargs=2, default=[2944, 1920],
  785. help='Crop size for images (height width)')
  786. parser.add_argument('--fold', type=int, default=1,
  787. help='which Fold to train on')
  789. # Wandb parameters
  790. parser.add_argument('--wandb-project', type=str, required=True,
  791. help='WandB project name')
  792. parser.add_argument('--wandb-run-name', type=str, required=True,
  793. help='WandB run name')
  795. args = parser.parse_args()
  796. train(args)

train.py at commit a1b230f, under CC-BY-NC-4.0 · at the source

Overview

Authors: Haresh Rengaraj Rajamohan1, Yanqi Xu1, Weicheng Zhu1, Richard Kijowski2, Kyunghyun Cho1, Krzysztof J. Geras1,3, Narges Razavian3,4, Cem M. Deniz3,5
ORCID iDs: Haresh Rengaraj Rajamohan, Yanqi Xu, Cem M. Deniz
  1. Center for Data Science, New York University, New York, New York, United States of America
  2. Department of Radiology, Hospital for Special Surgery, New York, New York, United States of America
  3. Department of Radiology, New York University Langone Health, New York, New York, United States of America
  4. Department of Population Health, New York University Langone Health, New York, New York, United States of America
  5. Bernard and Irene Schwartz Center for Biomedical Imaging, New York University Langone Health, New York, New York, United States of America
Institutions: New York University (United States); Hospital for Special Surgery (United States); NYU Langone Health (United States)
Journal: PloS one, volume 21, issue 5, article e0344600
Dates: received 26 September 2025; accepted 23 February 2026; published online 6 May 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1371/journal.pone.0344600 · PMID 42090385 · PMCID PMC13148697 · OpenAlex W4414497706
Open access: gold, a free copy (OpenAlex)
Preprint: doi:10.1101/2025.09.22.25336414
Status: code verified
Categories: human (organism), other condition (population), Alzheimer's / dementia (population), clinical / translational (subfield)
Methods: Machine learning, Statistics
MeSH: Alzheimer Disease*, Breast Neoplasms*, Deep Learning*, Osteoarthritis, Knee*, Disease Progression, Female, Humans, Magnetic Resonance Imaging, Predictive Learning Models, Prognosis, ROC Curve (* major topic)
Journal subjects: Medicine and Health Sciences, Diagnostic Medicine, Prognosis, Cancer Detection and Diagnosis, Oncology, Mental Health and Psychiatry, Dementia, Alzheimer's Disease, Neurology, Medical Conditions, Neurodegenerative Diseases, Rheumatology, Arthritis, Osteoarthritis, Biology and Life Sciences, Anatomy, Musculoskeletal System, Skeleton, Skeletal Joints, Knees, Body Limbs, Legs, Cancers and Neoplasms, Breast Tumors, Breast Cancer, Neuroscience, Cognitive Science, Cognitive Psychology, Learning, Psychology, Social Sciences, Learning and Memory
Topic: AI in cancer detection (Artificial Intelligence, Computer Science), according to OpenAlex
Funding: National Institute of Arthritis and Musculoskeletal and Skin Diseases (R01-AR074453, N/A (Salary))
Citations: cited by 1 paper (Europe PMC); 40 references in the paper

Abstract

Accurate disease prognosis is essential for patient care but is often hindered by the scarcity of longitudinal data. This study explores deep learning training strategies that utilize large, accessible diagnostic datasets to pretrain models aimed at predicting future disease progression in knee osteoarthritis (OA), Alzheimer’s disease (AD), and breast cancer (BC). While diagnostic pretraining improves prognostic task performance, naive fine-tuning for prognosis can cause ‘catastrophic forgetting,’ where the model’s original diagnostic accuracy degrades, a significant patient safety concern in real-world settings. To address this, we propose a sequential learning strategy with experience replay. We used cohorts with knee radiographs, brain MRIs, and digital mammograms to predict 4-year structural worsening in OA, 2-year cognitive decline in AD, and 5-year cancer diagnosis in BC. Our results showed that diagnostic pretraining on larger datasets improved prognosis model performance compared to standard baselines, boosting both the Area Under the Receiver Operating Characteristic curve (AUROC) (e.g., Knee OA external: 0.770 vs 0.747; Breast Cancer: 0.874 vs 0.848) and the Area Under the Precision-Recall Curve (AUPRC) (e.g., Alzheimer’s Disease: 0.752 vs 0.683). Additionally, a sequential learning approach with experience replay achieved prognostic performance comparable to dedicated single-task models (e.g., Breast Cancer AUROC 0.876 vs 0.874) while also preserving diagnostic ability. This method maintained high diagnostic accuracy (e.g., Breast Cancer Balanced Accuracy 50.4% vs 50.9% for a dedicated diagnostic model), unlike simpler multitask methods prone to catastrophic forgetting (e.g., 37.7%). Our findings show that leveraging large diagnostic datasets is a reliable and data-efficient way to enhance prognostic models while maintaining essential diagnostic skills.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repository

Its files are read in the Code ↔ Paper reader above, with 5 matches between paragraphs and lines of code.

denizlab/diag-to-prog-replay

License: CC-BY-NC-4.0
State: the link answers, verified on 28 September 2026
Evidence: files inventoried
Commit: a1b230ff75f5386020e6e1bce856dabd56d34b72, 7 January 2026
Languages: Python (51), Shell (6)
Size: 179 files, 57 scripts
Software Heritage: not archived
Found in: “Data Availability”
Holds: README, license file, environment (requirements.txt)
Not found: CITATION.cff, tests, continuous integration, documentation
Tools: PyTorch (44 files), NumPy (28 files), scikit-learn (11 files), pandas (9 files), Pillow (7 files), SciPy (6 files), h5py (5 files), FreeSurfer (4 files), Matplotlib (3 files), OpenCV (3 files), scikit-image (2 files), NiBabel (1 file)
Availability: 1 check, the latest on 28 September 2026: the link answers
  • 28 September 2026: the link answers
59 files
At the source: github.com/denizlab/diag-to-prog-replay

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 57 scripts, each with its path and the digest of its content;
  • 5 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Data Availability

- Osteoarthritis Initiative (OAI): OAI data are publicly available through the NIMH Data Archive (NDA). Access requires creating an NDA account, agreeing to the OAI data access terms, and requesting the desired collections through the NDA portal. Full instructions are provided at https://nda.nih.gov/oai. - Multicenter Osteoarthritis Study (MOST): MOST data are publicly available through the NIA Aging Research Biobank. Access requires creating a Biobank account and submitting a data request in accordance with the Biobank’s terms and conditions. More information is available at https://agingresearchbiobank.nia.nih.gov/. - Alzheimer’s Disease Neuroimaging Initiative (ADNI): ADNI data are publicly available from the LONI Image & Data Archive upon completion of a web application and acceptance of the ADNI Data Use Agreement. Data can be accessed at https://adni.loni.usc.edu/. - Breast Cancer Cohort (NYU Langone Health): This institutional dataset contains protected health information and cannot be made publicly available. De-identified data may be shared upon reasonable request, pending approval by NYU Langone’s data governance committee. Inquiries can be directed to Krzysztof Geras (). - Code and Model Availability: All code for preprocessing and model training sufficient to reproduce the analyses reported in this study are available at https://github.com/denizlab/diag-to-prog-replay.

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 28 September 2026: the first record

Recorded: type, language, journal, volume, issue, pages, dates, 8 authors, 11 MeSH terms, 1 funder, 35 references.

Cite

This paper

Rajamohan, H. R., Xu, Y., Zhu, W., Kijowski, R., Cho, K., Geras, K. J., Razavian, N., & Deniz, C. M. (2026). Robust disease prognosis via diagnostic knowledge preservation: A sequential learning approach. PloS one, 21(5), e0344600. https://doi.org/10.1371/journal.pone.0344600

BibTeX

@article{rajamohan2026robust,
author = {Rajamohan, Haresh Rengaraj and Xu, Yanqi and Zhu, Weicheng and Kijowski, Richard and Cho, Kyunghyun and Geras, Krzysztof J. and Razavian, Narges and Deniz, Cem M.},
title = {{Robust disease prognosis via diagnostic knowledge preservation: A sequential learning approach}},
journal = {PloS one},
year = {2026},
month = may,
volume = {21},
number = {5},
pages = {e0344600},
publisher = {PLOS},
issn = {1932-6203},
doi = {10.1371/journal.pone.0344600},
url = {https://doi.org/10.1371/journal.pone.0344600},
pmid = {42090385},
pmcid = {PMC13148697}
}

RIS

TY - JOUR
AU - Rajamohan, Haresh Rengaraj
AU - Xu, Yanqi
AU - Zhu, Weicheng
AU - Kijowski, Richard
AU - Cho, Kyunghyun
AU - Geras, Krzysztof J.
AU - Razavian, Narges
AU - Deniz, Cem M.
TI - Robust disease prognosis via diagnostic knowledge preservation: A sequential learning approach
T2 - PloS one
J2 - PLoS One
PY - 2026
DA - 2026/05/06
VL - 21
IS - 5
SP - e0344600
SN - 1932-6203
PB - PLOS
DO - 10.1371/journal.pone.0344600
UR - https://doi.org/10.1371/journal.pone.0344600
LA - en
ER -

CSL-JSON

{
"id": "10.1371/journal.pone.0344600",
"type": "article-journal",
"title": "Robust disease prognosis via diagnostic knowledge preservation: A sequential learning approach",
"container-title": "PloS one",
"author": [
{
"family": "Rajamohan",
"given": "Haresh Rengaraj"
},
{
"family": "Xu",
"given": "Yanqi"
},
{
"family": "Zhu",
"given": "Weicheng"
},
{
"family": "Kijowski",
"given": "Richard"
},
{
"family": "Cho",
"given": "Kyunghyun"
},
{
"family": "Geras",
"given": "Krzysztof J."
},
{
"family": "Razavian",
"given": "Narges"
},
{
"family": "Deniz",
"given": "Cem M."
}
],
"container-title-short": "PLoS One",
"volume": "21",
"issue": "5",
"page": "e0344600",
"DOI": "10.1371/journal.pone.0344600",
"PMID": "42090385",
"PMCID": "PMC13148697",
"ISSN": "1932-6203",
"publisher": "PLOS",
"URL": "https://doi.org/10.1371/journal.pone.0344600",
"language": "en",
"issued": {
"date-parts": [
[
2026,
5,
6
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.7554/elife.107933 [code]
Modality-agnostic decoding of vision and language from fMRI.
Journal: eLife
In common: FreeSurfer, OpenCV, scikit-image, 9 other tools, 1 reference
[2] doi:10.1016/j.patter.2026.101538 [code]
A multi-modal foundation model for brain disease diagnosis and medical imaging.
Journal: Patterns (New York, N.Y.)
In common: OpenCV, scikit-image, h5py, 8 other tools, clinical / translational, 1 reference
[3] doi:10.1038/s41597-026-07248-6 [code]
A large-scale fMRI dataset for vision-language semantic association.
Journal: Scientific data
In common: FreeSurfer, OpenCV, scikit-image, 9 other tools
[4] doi:10.1002/alz.71649 [code]
Postmortem brain MRI reveals differential associations of subcortical and limbic volumes with cortical thinning and neurodegenerative pathologies.
Journal: Alzheimer's & dementia : the journal of the Alzheimer's Association
In common: FreeSurfer, OpenCV, scikit-image, 8 other tools, Alzheimer's / dementia, other condition
[5] doi:10.1038/s41598-026-57519-w [code]
Automated segmentation of neurons and spinal cord structures in immunofluorescence images using SpineDL.
Journal: Scientific reports
In common: OpenCV, scikit-image, h5py, 8 other tools, other condition
[6] doi:10.1038/s41467-026-76098-y [code]
A single computational objective can produce specialization of streams in visual cortex.
Journal: Nature communications
In common: FreeSurfer, scikit-image, h5py, 8 other tools
[7] doi:10.1371/journal.pcbi.1014263 [code]
MIRAGE: Robust multi-modal architectures translate fMRI-to-image models from vision to mental imagery.
Journal: PLoS computational biology
In common: OpenCV, scikit-image, h5py, 8 other tools
[8] doi:10.3389/frai.2026.1771088 [code]
Few-shot deployment of pretrained MRI transformers in brain imaging tasks.
Journal: Frontiers in artificial intelligence
In common: OpenCV, scikit-image, h5py, 8 other tools
[9] doi:10.1038/s41467-026-73373-w [code]
Mapping neuro-vascular unit communications reveals distinct angiogenic programs across developing mouse brain regions.
Journal: Nature communications
In common: OpenCV, scikit-image, h5py, 8 other tools
[10] doi:10.1038/s41467-026-71555-0 [code]
A deep representation learning model to predict response to vagus nerve stimulation.
Journal: Nature communications
In common: FreeSurfer, NiBabel, PyTorch, 5 other tools, clinical / translational, 2 references

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.