Migrate Graph_TorchvisionModel to the current dataset API

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
 --------------
 

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",

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={},

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()