pyhealth.models.Agent#

The separate callable AgentLayer and the complete Agent model.

class pyhealth.models.AgentLayer(input_dim, static_dim=0, cell='gru', use_baseline=True, n_actions=10, n_units=64, n_hidden=128, dropout=0.5, lamda=0.5)[source]#

Bases: Module

Dr. Agent layer.

Paper: Junyi Gao et al. Dr. Agent: Clinical predictive model via mimicked second opinions. JAMIA.

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

Parameters:

  • input_dim (int) – dynamic feature size.

  • static_dim (int) – static feature size, if 0, then no static feature is used.

  • cell (str) – rnn cell type. Default is “gru”.

  • use_baseline (bool) – whether to use baseline for the RL agent. Default is True.

  • n_actions (int) – number of historical visits to choose. Default is 10.

  • n_units (int) – number of hidden units in each agent. Default is 64.

  • fusion_dim – number of hidden units in the final representation. Default is 128.

  • n_hidden (int) – number of hidden units in the rnn. Default is 128.

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

  • lamda (int) – weight for the agent selected hidden state and the current hidden state. Default is 0.5.

Examples

>>> from pyhealth.models import AgentLayer
>>> input = torch.randn(3, 128, 64)  # [batch size, sequence len, feature_size]
>>> layer = AgentLayer(64)
>>> c, _ = layer(input)
>>> c.shape
torch.Size([3, 128])
choose_action(observation, agent=1)[source]#
forward(x, static=None, mask=None)[source]#

Forward propagation.

Parameters:

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

  • static (Optional[tensor]) – a tensor of shape [batch size, static_dim].

  • 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, n_hidden] representing the

patient embedding.

output: a tensor of shape [batch size, sequence len, n_hidden] representing the patient embedding at each time step.

Return type:

last_output

training: bool#
class pyhealth.models.Agent(dataset, feature_keys, label_key, mode, static_key=None, embedding_dim=128, hidden_dim=128, use_baseline=True, **kwargs)[source]#

Bases: BaseModel

Dr. Agent model.

Paper: Junyi Gao et al. Dr. Agent: Clinical predictive model via mimicked second opinions. JAMIA.

Note

We use separate Dr. Agent layers for different feature_keys. Currently, 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 Dr. Agent 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 Dr. Agent 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 Dr. Agent 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 Dr. Agent 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 Dr. Agent 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”.

  • static_keys – the key in samples to use as static features, e.g. “demographics”. Default is None. we only support numerical static features.

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

  • hidden_dim (int) – the hidden dimension of the RNN in the Dr. Agent layer. Default is 128.

  • use_baseline (bool) – whether to use the baseline value to calculate the RL loss. Default is True.

  • **kwargs – other parameters for the Dr. Agent 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]],
...             ],
...             "demographic": [0.0, 2.0, 1.5],
...             "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]],
...             ],
...             "demographic": [0.0, 2.0, 1.5],
...             "label": 0,
...         },
...     ]
>>> dataset = SampleEHRDataset(samples=samples, dataset_name="test")
>>>
>>> from pyhealth.models import Agent
>>> model = Agent(
...         dataset=dataset,
...         feature_keys=[
...             "list_codes",
...             "list_vectors",
...             "list_list_codes",
...             "list_list_vectors",
...         ],
...         label_key="label",
...         static_key="demographic",
...         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(1.4059, grad_fn=<AddBackward0>),
    'y_prob': tensor([[0.4861], [0.5348]], grad_fn=<SigmoidBackward0>),
    'y_true': tensor([[0.], [1.]]),
    'logit': tensor([[-0.0556], [0.1392]], grad_fn=<AddmmBackward0>)
}
>>>
get_loss(model, pred, true, mask, gamma=0.9, entropy_term=0.01)[source]#
forward(**kwargs)[source]#

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 final loss. loss_task: a scalar tensor representing the task loss. loss_RL: a scalar tensor representing the RL 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#