indiv_outcome.sequence.RETAIN

class pytrial.tasks.indiv_outcome.sequence.retain.RETAIN(vocab_size, orders, mode, output_dim=None, max_visit=None, emb_size=64, n_rnn_layer=2, learning_rate=0.0001, weight_decay=0.0001, batch_size=64, epochs=10, num_worker=0, device='cuda:0', experiment_id='test')

Implement RETAIN for longitudinal patient records predictive modeling 1.

Parameters

• vocab_size (list[int]) – A list of vocabulary size for different types of events, e.g., for diagnosis, procedure, medication.

• orders (list[str]) – A list of orders when treating inputs events. Should have the same shape of vocab_size.

• output_dim (int) –

  The dimension of the output.

  • If binary classification, output_dim=1;

  • If multiclass/multilabel classification, output_dim=n_class

  • If regression, output_dim=1.

• mode (str) – Prediction traget in [‘binary’,’multiclass’,’multilabel’,’regression’].

• max_visit (int) – The maximum number of visits for input event codes.

• emb_size (int) – Embedding size for encoding input event codes.

• n_rnn_layer (int) – Number of RETAIN layers for encoding historical events.

• learning_rate (float) – Learning rate for optimization based on SGD. Use torch.optim.Adam by default.

• weight_decay (float) – Regularization strength for l2 norm; must be a positive float. Smaller values specify weaker regularization.

• batch_size (int) – Batch size when doing SGD optimization.

• epochs (int) – Maximum number of iterations taken for the solvers to converge.

• num_worker (int) – Number of workers used to do dataloading during training.

• device (str) – The model device.

Notes

1

Choi, E., Bahadori, M. T., Sun, J., Kulas, J., Schuetz, A., & Stewart, W. (2016). Retain: An interpretable predictive model for healthcare using reverse time attention mechanism. Advances in neural information processing systems, 29.

fit(train_data, valid_data=None)

Train model with sequential patient records.

Parameters

• train_data (SequencePatientBase) – A SequencePatientBase contains patient records where ‘v’ corresponds to visit sequence of different events; ‘y’ corresponds to labels.

• valid_data (SequencePatientBase) – A SequencePatientBase contains patient records used to make early stopping of the model.

load_model(checkpoint)

Load pretrained model from the disk.

Parameters

checkpoint (str) – The input directory that stores the trained pytorch model and configuration.

predict(test_data)

Predict patient outcomes using longitudinal trial patient sequences.

Parameters

test_data (SequencePatient) – A SequencePatient contains patient records where ‘v’ corresponds to visit sequence of different events.

save_model(output_dir)

Save the pretrained model to the disk.

Parameters

output_dir (str) – The output directory that stores the trained pytorch model and configuration.