diff --git a/docs/api/models.rst b/docs/api/models.rst index 4187c123b..b1b360054 100644 --- a/docs/api/models.rst +++ b/docs/api/models.rst @@ -74,6 +74,15 @@ varies by task type: The ``forward()`` method is expected to return a dictionary with four keys: ``loss``, ``y_prob``, ``y_true``, and ``logit``. The Trainer reads all four. +Graph Image Models +------------------ + +``Graph_TorchvisionModel`` takes a ``SampleDataset``, ``model_name``, +``model_config``, and ``gnn_config``. Its single image field, label field, and +classification mode come from the dataset schemas. Pass processed image and +label tensors from ``get_dataloader`` to ``forward``, together with two +``EdgeIndex`` adjacencies under ``adjacencies``, aligned with the image batch. + EmbeddingModel -------------- diff --git a/examples/graph_torchvision_model.ipynb b/examples/graph_torchvision_model.ipynb index faaf7cf9a..45dbe5747 100644 --- a/examples/graph_torchvision_model.ipynb +++ b/examples/graph_torchvision_model.ipynb @@ -190,9 +190,6 @@ "\n", "model = Graph_TorchvisionModel(\n", " dataset=sample_dataset,\n", - " feature_keys=[\"path\"],\n", - " label_key=\"label\",\n", - " mode=\"multiclass\",\n", " model_name=\"resnet18\",\n", " model_config={},\n", " gnn_config={\"input_dim\": 256, \"hidden_dim\": 128},\n", diff --git a/pyhealth/models/graph_torchvision_model.py b/pyhealth/models/graph_torchvision_model.py index f91350cee..68f46f567 100644 --- a/pyhealth/models/graph_torchvision_model.py +++ b/pyhealth/models/graph_torchvision_model.py @@ -7,7 +7,6 @@ import math import sys -from typing import Dict, List import torch import torch.nn as nn @@ -112,6 +111,13 @@ def __repr__(self): class GCN(nn.Module): + """Two-layer graph network returning unnormalized class logits. + + Example: + >>> gcn = GCN(8, 4, 2, dropout=0.5, init="xavier") + >>> logits = gcn(node_features, adjacencies) + """ + def __init__(self, nfeat, nhid, nclass, dropout, init): super(GCN, self).__init__() @@ -132,7 +138,7 @@ def forward(self, x, adjs): temp = self.to_sparse_adj(adjs[1], size = (adjs[0].size[0], adjs[0].size[0])) x = self.gc2(x, temp) - return F.log_softmax(x, dim=1) + return x class Graph_TorchvisionModel(BaseModel): @@ -168,42 +174,42 @@ class Graph_TorchvisionModel(BaseModel): ----------------------------------------------------------------------------------- Args: - dataset: the dataset to train the model. It is used to query certain - information such as the set of all tokens. - feature_keys: list of keys in samples to use as features, e.g., ["image"]. - Only one feature is supported. - label_key: key in samples to use as label, e.g., "drugs". - mode: one of "binary", "multiclass", or "multilabel". + dataset: dataset with one image input and one binary, multiclass, or + multilabel output. Field names and mode are derived from its schema. model_name: str, name of the model to use, e.g., "resnet18". See SUPPORTED_MODELS in the source code for the full list. model_config: dict, kwargs to pass to the model constructor, e.g., {"weights": "DEFAULT"}. See the torchvision documentation for the set of supported kwargs for each model. + gnn_config: graph network dimensions, with keys "input_dim" and "hidden_dim". + + Example: + >>> model = Graph_TorchvisionModel( + ... dataset, model_name="resnet18", model_config={"weights": None}, + ... gnn_config={"input_dim": 256, "hidden_dim": 128}, + ... ) + >>> result = model(**batch, adjacencies=adjacencies) + >>> result["loss"].backward() ----------------------------------------------------------------------------------- """ def __init__( self, dataset: SampleDataset, - feature_keys: List[str], - label_key: str, - mode: str, model_name: str, model_config: dict, gnn_config: dict, ): - super(Graph_TorchvisionModel, self).__init__( - dataset=dataset, - feature_keys=feature_keys, - label_key=label_key, - mode=mode, - ) + super().__init__(dataset=dataset) self.model_name = model_name self.model_config = model_config self.gnn_config = gnn_config - assert len(feature_keys) == 1, "Only one feature is supported!" + assert len(self.feature_keys) == 1, "Only one image feature is supported!" + assert len(self.label_keys) == 1, "Only one label is supported!" + self.feature_key = self.feature_keys[0] + self.label_key = self.label_keys[0] assert model_name in SUPPORTED_MODELS_FINAL_LAYER.keys(), \ f"PyHealth does not currently include {model_name} model!" @@ -219,14 +225,13 @@ def __init__( gnn_input_dim = gnn_config["input_dim"] gnn_hidden_dim = gnn_config["hidden_dim"] - self.label_tokenizer = self.get_label_tokenizer() - output_size = self.get_output_size(self.label_tokenizer) + output_size = self.get_output_size() self.gnn = GCN(nfeat=gnn_input_dim, nhid=gnn_hidden_dim, nclass=output_size, dropout=0.5, init='uniform') setattr(self.model, final_layer_name.split(".")[0], nn.Linear(hidden_dim, gnn_input_dim)) - def build_graph(self, data, random = False) -> Dict[str, torch.Tensor]: + def build_graph(self, data, random = False) -> dict[str, torch.Tensor]: """This module generate edge index of graph structure based on given data. Currently, we do not have multi-modal data, so this module randomly generate edge index""" @@ -238,20 +243,19 @@ def build_graph(self, data, random = False) -> Dict[str, torch.Tensor]: } - def forward(self, **kwargs) -> Dict[str, torch.Tensor]: + def forward(self, **kwargs) -> dict[str, torch.Tensor]: """Forward propagation.""" - # concat the info within one batch (batch, channel, length) - x = kwargs["image"] - x = torch.stack(x, dim=0).to(self.device) + x = kwargs[self.feature_key].to(self.device) if x.shape[1] == 1: x = x.repeat((1, 3, 1, 1)) img_embs = self.model(x) logits = self.gnn(img_embs, kwargs["adjacencies"]) - y_true = self.prepare_labels(kwargs[self.label_key], self.label_tokenizer) + y_true = kwargs[self.label_key].to(self.device) loss = self.get_loss_function()(logits, y_true) y_prob = self.prepare_y_prob(logits) return { "loss": loss, + "logit": logits, "y_prob": y_prob, "y_true": y_true, } @@ -293,9 +297,6 @@ def encode(sample): model = Graph_TorchvisionModel( dataset=sample_dataset, - feature_keys=["path"], - label_key="label", - mode="multiclass", model_name="resnet18", # model_config={"weights": "DEFAULT"}, model_config={}, diff --git a/tests/core/test_graph_torchvision_model.py b/tests/core/test_graph_torchvision_model.py new file mode 100644 index 000000000..2c722d437 --- /dev/null +++ b/tests/core/test_graph_torchvision_model.py @@ -0,0 +1,63 @@ +import pytest +import torch +from PIL import Image + +from pyhealth.datasets import create_sample_dataset, get_dataloader +from pyhealth.models import Graph_TorchvisionModel +from pyhealth.processors import ImageProcessor, MultiClassLabelProcessor +from pyhealth.sampler.sage_sampler import EdgeIndex + + +@pytest.mark.parametrize( + "processor,labels,output_size,mode", + [ + ("binary", [0, 1], 1, "L"), + (MultiClassLabelProcessor, ["healthy", "ill"], 2, "RGB"), + ("multilabel", [["a"], ["b"]], 2, "L"), + ], +) +def test_graph_torchvision_processed_batch(tmp_path, processor, labels, output_size, mode): + path = tmp_path / "scan.png" + Image.new(mode, (32, 32), color=128).save(path) + dataset = create_sample_dataset( + samples=[{"scan": str(path), "target": label} for label in labels], + input_schema={"scan": ImageProcessor(image_size=32, mode=mode)}, + output_schema={"target": processor}, + ) + model = Graph_TorchvisionModel( + dataset, + model_name="resnet18", + model_config={"weights": None}, + gnn_config={"input_dim": 8, "hidden_dim": 4}, + ) + batch = next(iter(get_dataloader(dataset, batch_size=2))) + adjacency = EdgeIndex(torch.tensor([[0, 1], [0, 1]]), None, (2, 2)) + # Keep the classifier deterministic and away from inactive ReLUs. + model.eval() + with torch.no_grad(): + model.gnn.gc1.bias.fill_(1) + model.gnn.gc2.weight.fill_(0.1) + model.gnn.gc2.bias.fill_(0.2) + result = model(**batch, adjacencies=[adjacency, adjacency]) + + assert model.feature_keys == ["scan"] + assert model.label_keys == ["target"] + assert result["y_prob"].shape == (2, output_size) + torch.testing.assert_close(result["y_true"], batch["target"]) + torch.testing.assert_close( + result["y_prob"], model.prepare_y_prob(result["logit"]) + ) + assert torch.isfinite(result["loss"]) + result["loss"].backward() + for parameter in (model.model.fc.weight, model.gnn.gc2.weight): + assert parameter.grad is not None + assert torch.isfinite(parameter.grad).all() + if processor == "binary": + assert not torch.all(result["y_prob"] == 0.5) + assert model.gnn.gc2.weight.grad.abs().sum() > 0 + + model.train() + model.zero_grad() + loss = model(**batch, adjacencies=[adjacency, adjacency])["loss"] + assert torch.isfinite(loss) + loss.backward()