diff --git a/pyhealth/models/biot.py b/pyhealth/models/biot.py index e038e092b..d3b4b00f0 100644 --- a/pyhealth/models/biot.py +++ b/pyhealth/models/biot.py @@ -129,6 +129,18 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: class BIOTEncoder(nn.Module): + """Encoder for multichannel biosignals. + + Examples: + >>> _ = _get_linear_attention_transformer() + >>> encoder = BIOTEncoder( + ... emb_size=128, heads=8, depth=2, n_channels=18 + ... ) + >>> signal = torch.randn(2, 18, 2000) + >>> encoder(signal).shape + torch.Size([2, 128]) + """ + def __init__( self, emb_size=256, @@ -158,7 +170,7 @@ def __init__( self.positional_encoding = PositionalEncoding(emb_size) # channel token, N_channels >= your actual channels - self.channel_tokens = nn.Embedding(n_channels, 256) + self.channel_tokens = nn.Embedding(n_channels, emb_size) self.index = nn.Parameter( torch.LongTensor(range(n_channels)), requires_grad=False ) @@ -372,4 +384,4 @@ def get_embeddings(self, **kwargs: Any) -> Dict[str, torch.Tensor]: print(f"āœ“ BIOTClassifier forward pass:") print(f" Logits shape: {logits.shape}") - print("\nāœ“ All tests passed!") \ No newline at end of file + print("\nāœ“ All tests passed!") diff --git a/tests/core/test_biot.py b/tests/core/test_biot.py index d51eecb41..32836dcb3 100644 --- a/tests/core/test_biot.py +++ b/tests/core/test_biot.py @@ -176,6 +176,32 @@ def test_model_different_n_classes(self): self.assertEqual(ret["logit"].shape[1], 1) + def test_model_non_default_emb_size(self): + """BIOT must honor a non-default emb_size. + + Regression test: the channel-token embedding dimension was hardcoded + to 256, so it could not be added to the emb_size-dimensional spectral + embedding when emb_size != 256, crashing the forward pass. + """ + model = BIOT( + dataset=self.dataset, + emb_size=128, + heads=8, + depth=2, + n_fft=200, + hop_length=100, + n_channels=18, + ) + self.assertEqual(model.biot.biot.channel_tokens.weight.shape[1], 128) + + train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=False) + data_batch = next(iter(train_loader)) + ret = model(**data_batch) + ret["loss"].backward() + + expected_size = self.dataset.output_processors["label"].size() + self.assertEqual(ret["logit"].shape[1], expected_size) + def test_model(self): """Test BIOT""" model_small = BIOT(