pyhealth.models.RETAIN

The separate callable RETAINLayer and the complete RETAIN model.

class pyhealth.models.RETAINLayer(feature_size, dropout=0.5)

Bases: Module

RETAIN layer.

Paper: Edward Choi et al. RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism. NIPS 2016.

This layer is used in the RETAIN model. But it can also be used as a standalone layer.

Parameters:

• feature_size (int) – the hidden feature size.

• dropout (float) – dropout rate. Default is 0.5.

Examples

>>> from pyhealth.models import RETAINLayer
>>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
>>> layer = RETAINLayer(64)
>>> c = layer(input)
>>> c.shape
torch.Size([3, 64])

static reverse_x(input, lengths)

Reverses the input.

compute_alpha(rx, lengths)

Computes alpha attention.

compute_beta(rx, lengths)

Computes beta attention.

forward(x, mask=None)

Forward propagation.

Parameters:

• x (tensor) – a tensor of shape [batch size, sequence len, feature_size].

• mask (Optional[tensor]) – an optional tensor of shape [batch size, sequence len], where 1 indicates valid and 0 indicates invalid.

Returns:

a tensor of shape [batch size, feature_size] representing the

context vector.

Return type:

c

training: bool

class pyhealth.models.RETAIN(dataset, feature_keys, label_key, mode, pretrained_emb=None, embedding_dim=128, **kwargs)

Bases: BaseModel

RETAIN model.

Paper: Edward Choi et al. RETAIN: An Interpretable Predictive Model for Healthcare using Reverse Time Attention Mechanism. NIPS 2016.

Note

We use separate Retain layers for different feature_keys. Currentluy, we automatically support different input formats:

• code based input (need to use the embedding table later)

• float/int based value input

We follow the current convention for the Retain model:

• case 1. [code1, code2, code3, …]

  • we will assume the code follows the order; our model will encode

  each code into a vector and apply Retain on the code level

• case 2. [[code1, code2]] or [[code1, code2], [code3, code4, code5], …]

  • we will assume the inner bracket follows the order; our model first

  use the embedding table to encode each code into a vector and then use average/mean pooling to get one vector for one inner bracket; then use Retain one the braket level

• case 3. [[1.5, 2.0, 0.0]] or [[1.5, 2.0, 0.0], [8, 1.2, 4.5], …]

  • this case only makes sense when each inner bracket has the same length;

  we assume each dimension has the same meaning; we run Retain directly on the inner bracket level, similar to case 1 after embedding table

• case 4. [[[1.5, 2.0, 0.0]]] or [[[1.5, 2.0, 0.0], [8, 1.2, 4.5]], …]

  • this case only makes sense when each inner bracket has the same length;

  we assume each dimension has the same meaning; we run Retain directly on the inner bracket level, similar to case 2 after embedding table

Parameters:

• dataset (SampleEHRDataset) – the dataset to train the model. It is used to query certain information such as the set of all tokens.

• feature_keys (List[str]) – list of keys in samples to use as features, e.g. [“conditions”, “procedures”].

• label_key (str) – key in samples to use as label (e.g., “drugs”).

• mode (str) – one of “binary”, “multiclass”, or “multilabel”.

• embedding_dim (int) – the embedding dimension. Default is 128.

• **kwargs – other parameters for the RETAIN layer.

Examples

>>> from pyhealth.datasets import SampleEHRDataset
>>> samples = [
...         {
...             "patient_id": "patient-0",
...             "visit_id": "visit-0",
...             "list_codes": ["505800458", "50580045810", "50580045811"],  # NDC
...             "list_vectors": [[1.0, 2.55, 3.4], [4.1, 5.5, 6.0]],
...             "list_list_codes": [["A05B", "A05C", "A06A"], ["A11D", "A11E"]],  # ATC-4
...             "list_list_vectors": [
...                 [[1.8, 2.25, 3.41], [4.50, 5.9, 6.0]],
...                 [[7.7, 8.5, 9.4]],
...             ],
...             "label": 1,
...         },
...         {
...             "patient_id": "patient-0",
...             "visit_id": "visit-1",
...             "list_codes": [
...                 "55154191800",
...                 "551541928",
...                 "55154192800",
...                 "705182798",
...                 "70518279800",
...             ],
...             "list_vectors": [[1.4, 3.2, 3.5], [4.1, 5.9, 1.7], [4.5, 5.9, 1.7]],
...             "list_list_codes": [["A04A", "B035", "C129"]],
...             "list_list_vectors": [
...                 [[1.0, 2.8, 3.3], [4.9, 5.0, 6.6], [7.7, 8.4, 1.3], [7.7, 8.4, 1.3]],
...             ],
...             "label": 0,
...         },
...     ]
>>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
>>>
>>> from pyhealth.models import RETAIN
>>> model = RETAIN(
...         dataset=dataset,
...         feature_keys=[
...             "list_codes",
...             "list_vectors",
...             "list_list_codes",
...             "list_list_vectors",
...         ],
...         label_key="label",
...         mode="binary",
...     )
>>>
>>> from pyhealth.datasets import get_dataloader
>>> train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
>>> data_batch = next(iter(train_loader))
>>>
>>> ret = model(**data_batch)
>>> print(ret)
{
    'loss': tensor(0.5640, grad_fn=<BinaryCrossEntropyWithLogitsBackward0>),
    'y_prob': tensor([[0.5325],
                    [0.3922]], grad_fn=<SigmoidBackward0>),
    'y_true': tensor([[1.],
                    [0.]]),
    'logit': tensor([[ 0.1303],
                    [-0.4382]], grad_fn=<AddmmBackward0>)
}
>>>

forward(**kwargs)

Forward propagation.

The label kwargs[self.label_key] is a list of labels for each patient.

Parameters:

**kwargs – keyword arguments for the model. The keys must contain all the feature keys and the label key.

Returns:

loss: a scalar tensor representing the loss. y_prob: a tensor representing the predicted probabilities. y_true: a tensor representing the true labels.

Return type:

A dictionary with the following keys

training: bool