rename

Author: Logiquo

pyhealth/models/stagenet_mha.py

@@ -275,11 +275,12 @@ def forward(
         return last_output, output, distance
 
 
-class StageNet(BaseModel):
+class StageAttentionNet(BaseModel):
     """StageNet model.
 
     Paper: Junyi Gao et al. Stagenet: Stage-aware neural networks for health
-    risk prediction. WWW 2020.
+    risk prediction. WWW 2020. But with Multi-Head Attention (MHA) between
+    the SA-LSTM and the SA-CNN.
 
     This model uses the StageNetProcessor which expects inputs in the format:
         {"value": [...], "time": [...]}
@@ -376,7 +377,7 @@ def __init__(
         levels: int = 3,
         **kwargs,
     ):
-        super(StageNet, self).__init__(
+        super(StageAttentionNet, self).__init__(
             dataset=dataset,
         )
         self.embedding_dim = embedding_dim
@@ -689,7 +690,7 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
     train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
 
     # model
-    model = StageNet(dataset=dataset)
+    model = StageAttentionNet(dataset=dataset)
 
     # data batch
     data_batch = next(iter(train_loader))

tests/core/test_stagenet_mha.py

@@ -2,7 +2,7 @@
 import torch
 
 from pyhealth.datasets import SampleDataset, get_dataloader
-from pyhealth.models.stagenet_mha import StageNet as StageNetMHA
+from pyhealth.models.stagenet_mha import StageAttentionNet as StageNetMHA
 
 
 class TestStageNetMHA(unittest.TestCase):