Function main

train.py:305–409

(args)

def main(args):

    # Set random seed for some degree of reproducibility. See PyTorch docs on this topic for caveats.
    # https://pytorch.org/docs/stable/notes/randomness.html#reproducibility
    set_seed()

    if not torch.cuda.is_available():
        raise Exception(
            "No CUDA device available. Training without one is not feasible."
        )

    git_sha = subprocess.check_output(["git", "describe", "--always"]).strip().decode("utf-8")
    train_run_id = f"{git_sha}_hp{args.hp}_{time.strftime('%Y%m%d-%H%M')}"

    train_split, val_split, train_BS, test_BS, bins_values, stage_taxonomy = prepare_datasets(args)

    print(f"=> Run ID {train_run_id}")
    print(f"=> Training on {len(train_split)} samples")
    print(f"=> Validating on {len(val_split)} samples")

    base_hparams = dict( 
        # Preprocessing settings. This should be changed with the dataset called accordingly.
        # Storing values here for readibility.
        n_bins=args.n_bins, # Partion on the continuous time scale.
        bins_values=bins_values,
        input_feature_size=args.input_feature_size,
        features_extraction=os.path.dirname(args.data_dir),

        # Settings that be changed in the loop:
        # Training.
        sampling_method="random",
        max_epochs=100,
        earlystop_warmup=0,
        earlystop_patience=30,
        earlystop_min_epochs=30,

        # Loss.
        alpha_surv = 0.0,

        # Optimizer.
        initial_lr=0.00003,
        milestones="2, 5, 15, 25",
        gamma_lr=0.1,
        weight_decay=0.00001,

        # Model architecture parameters. See model class for details.
        precompression_layer=True,
        feature_size_comp=512,
        feature_size_attn=256,
        postcompression_layer=True,
        feature_size_comp_post=128,
        p_dropout_fc=0.25,
        p_dropout_atn=0.25,

        # Model of molecular classification. In our case only inference is used. 
        n_classes_molecular=args.n_classes_molecular,
        feature_size_comp_molecular=args.feature_size_comp_molecular,
        feature_size_attn_molecular=args.feature_size_attn_molecular,