Add switchable DIFF V1 attention

train_all_future.py

@@ -25,6 +25,7 @@ from torch.optim import AdamW
 from torch.utils.data import DataLoader, RandomSampler
 from tqdm.auto import tqdm
 
+from attention_types import DEFAULT_ATTENTION_TYPE, SUPPORTED_ATTENTION_TYPES
 from dataset import (
     DISEASE_HISTORY_MODES,
     DISEASE_HISTORY_MODE_TIMED,
@@ -119,6 +120,12 @@ def parse_args() -> argparse.Namespace:
         default=DEFAULT_MODEL_ARCHITECTURE,
         choices=SUPPORTED_MODEL_ARCHITECTURES,
     )
+    parser.add_argument(
+        "--attention_type",
+        type=str,
+        default=DEFAULT_ATTENTION_TYPE,
+        choices=SUPPORTED_ATTENTION_TYPES,
+    )
 
     parser.add_argument("--batch_size", type=int, default=128)
     parser.add_argument("--base_lr", type=float, default=3e-4)
@@ -204,6 +211,7 @@ def build_model(
         dist_mode=args.dist_mode,
         dropout=args.dropout,
         model_architecture=args.model_architecture,
+        attention_type=args.attention_type,
     )
 
 
@@ -333,6 +341,7 @@ def build_metadata(
         "collate_fn": "all_future_collate_fn",
         "model_class": "DeepHealth",
         "model_architecture": args.model_architecture,
+        "attention_type": args.attention_type,
         "model_target_mode": "all_future",
         "target_mode": "all_future",
         "event_stream_version": "disease_death_only_v1",
@@ -393,6 +402,7 @@ def main() -> None:
     logger.info(f"Starting all-future training run: {run_name}")
     logger.info(f"Device: {device}")
     logger.info(f"Model architecture: {args.model_architecture}")
+    logger.info(f"Attention type: {args.attention_type}")
     logger.info(f"Disease history mode: {args.disease_history_mode}")
     logger.info(f"extra_info_types: {format_extra_info_types(args.extra_info_types)}")
     logger.info("Continuous value scaling: RobustScale (required)")