Help with modelling a multi-input trainer.

[ak3ra]

Help, I am attempting to create a Cutom task using the templates provided;
the inputs are (in batches of two) :

sentinel2 rasters of size (2, 13, 600, 600)
esri landcover rasters of size (2, 1, 600, 600)
and the target (y) is a 100 class probability vector:

Here is my class:

class CustomModelTask(ClassificationTask):
  def config_task(self) -> None:
    self.config_model()

    if self.hyperparams["loss"] == "bce":
      self.loss = nn.BCEWithLogitsLoss()
    else:
      raise ValueError(f"loss type '{self.hyperparams['loss']}' is not valid.")

  def __init__(self, **kwargs: Any) -> None:
    super().__init__(**kwargs)

    self.input_features = 100
    self.output_features = 100
    self.classification_head = {"LAYER_COUNT": 3}

    self.save_hyperparameters()  # type: ignore[operator]
    self.hyperparams = cast(Dict[str, Any], self.hparams)
    self.config_task()
    self.backbone = resnet50(pretrained=False)

  def forward(self,x):
    self.backbone.conv1 = torch.nn.Conv2d(in_channels=x.shape[1], out_channels = 64,
                                          kernel_size=(7,7),
                                          stride=(2,2),
                                          bias=False)
    self.backbone.fc = torch.nn.Linear(self.backbone.fc.in_features,100)
    
    self.fc = torch.nn.Sequential(
                *([
                    torch.nn.Sequential(
                        torch.nn.Linear(in_features=self.input_features,out_features=self.output_features),
                        torch.nn.ReLU()
                    )
                    for _ in range(self.classification_head.get("LAYER_COUNT", 5) - 1)
                ])
            )
    
    
    output = self.backbone(x)
    
    self.label_prediction = torch.nn.Linear(400, 100)
    self.label_count = torch.nn.Linear(400, 100)
    
    return output

  def training_step(self, *args: Any, **kwargs: Any)-> Tensor:
    batch = args[0]
    x1 = batch["s2_image"]
    x2 = batch["lc_image"]
    y = batch["label"]
    output1 = self.forward(x1)
    # output2 = self.forward(x2)
    output1 = output1.flatten()
    
    x2_flat = x2.flatten()
    combined = torch.cat((output1, x2_flat), dim=0).to(torch.float32)
    outputs = self.fc(combined)
    y_hat  = self.label_prediction(outputs)
    label_count  = self.label_count(outputs)
    # TODO: Add loss function and everything else:
   
    y_hat = y_hat.unsqueeze(0)
    loss = self.loss(y_hat, y)
    print(loss)
    self.log("train_loss", loss, on_step = True, on_epoch=False)
    self.train_metrics(y_hat, y)
    
    return cast(Tensor, loss)

I however run into mismatch errors of the size of the target vector vs the predicted vector;

Question:

1. What is the best practice for using a convnet that can take multiple inputs as described above, using the formats of the tasks (RegressionTask and ClassificationTask )
2. What is the right way to fuse the multiple inputs to collectively learn a target?

[adamjstewart]

Personally, I would concatenate the (2, 13, 600, 600) Sentinel image with the (2, 1, 600, 600) esri land cover mask to get a single (2, 14, 600, 600) input image. Then I would pass this through a model like normal. So first create a dataset that does the concatenation, then use the normal builtin ClassificationTask trainer.

[isaaccorley]

To add onto this, as @adamjstewart mentioned, it's likely easier to create a dataset and/or datamodule that concatenates the S2 image and land cover mask and define the trainer like below for example:

dataset = S2ESRIConcatDataset(...)
dataloader = torch.utils.data.DataLoader(dataset, ...)
task = ClassificationTask(model="resnet18", weights="imagenet", in_channels=15, num_classes=100)
trainer = pl.Trainer(...)
trainer.fit(model=task, train_dataloaders=dataloader) # or if using a datamodule trainer.fit(model=task, datamodule=datamodule) 

However, if you want a multi output classifier and regression model you will need to create your own custom trainer. An example could be something like below. Of course there is much flexibility on what metrics you should use

import torch.nn as nn
import timm
from torchgeo.trainers import ClassificationTask

class CustomModelTask(ClassificationTask):
  def config_task(self) -> None:
    self.config_model()
    self.cls_loss_fn = nn.BCEWithLogitsLoss()
    self.reg_loss_fn = nn.MSELoss()

  def config_model(self) -> None:
    num_features = 512 # 512 for resnet18
    self.model = timm.create_model(
        self.hparams["model"],
        num_classes=0,
        in_chans=self.hparams["in_channels"],
        pretrained=self.hparams["weights"],
    )
   cls_head = nn.Linear(num_features, self.hparams["num_classes"])
   count_head = nn.Linear(num_features, self.hparams["num_classes"])

    def __init__(self, **kwargs: Any) -> None:
        super().__init__()
        self.save_hyperparameters()
        self.hyperparams = cast(Dict[str, Any], self.hparams)
        self.config_task()

        self.train_metrics_cls = MetricCollection(
            {
                "OverallAccuracy": MulticlassAccuracy(
                    num_classes=self.hyperparams["num_classes"], average="micro"
                ),
                "AverageAccuracy": MulticlassAccuracy(
                    num_classes=self.hyperparams["num_classes"], average="macro"
                ),
                "JaccardIndex": MulticlassJaccardIndex(
                    num_classes=self.hyperparams["num_classes"]
                ),
                "F1Score": MulticlassFBetaScore(
                    num_classes=self.hyperparams["num_classes"],
                    beta=1.0,
                    average="micro",
                ),
            },
            prefix="train_",
        )
        self.val_metrics_cls = self.train_metrics_cls.clone(prefix="val_")
        self.test_metrics_cls = self.train_metrics_cls.clone(prefix="test_")

        self.train_metrics_reg = MetricCollection(
            {"RMSE": MeanSquaredError(squared=False), "MAE": MeanAbsoluteError()},
            prefix="train_",
        )
        self.val_metrics_reg = self.train_metrics_reg.clone(prefix="val_")
        self.test_metrics_reg = self.train_metrics_reg.clone(prefix="test_")

  def forward(self,x):
    features = self.model(x)
    y_hat_cls = self.cls_head(features)
    y_hat_count = self.count_head(features)
    return (y_hat_cls, y_hat_count)

  def training_step(self, *args: Any, **kwargs: Any)-> Tensor:
    batch = args[0]
    x, y_cls, y_count = batch["image"], batch["label_cls"], batch["label_count"]
    y_hat_cls, y_hat_count = self(x)
    loss_cls = self.cls_loss_fn(y_hat_cls, y_cls)
    loss_reg = self.reg_loss_fn(y_hat_count, y_count)
    loss = loss_cls + loss_reg # this can be a weighted sum however you choose
    self.log("train_loss", loss, on_step = True, on_epoch=False)
    self.train_metrics_cls(y_hat_cls, y_cls)
    self.train_metrics_reg(y_hat_count, y_count)
    return cast(Tensor, loss)

[ak3ra]

Perfect!

Many thanks, @adamjstewart and @isaaccorley, I have incorporated both and it's working as expected!