Repository: ludwig-ai/ludwig
Pull request: Fix TorchVision channel preprocessing
This PR introduces automatic channel resizing for torchvision models to ensure that images fed into torchvision encoders always have 3 channels. Closes #3170.

Problem statement:
Clarification about TorchVision Pretrained Model Encoders usage
Hi,

This is not really a bug but more like a question for clarification since I might not understand the documentation properly.

What I want to do is to test various encoder architectures for my image classification problem. So I was quite happy when I saw the TorchVision encoder zoo popping up in the newest release and I was hoping to have access to many more architectures compared to `v0.6` and older. I now wanted to try the TorchVision Pretrained Model Encoders (without loading the weights) and train from scratch. E.g. just utilizing the architecture since my data is not very close to ImageNet content at all. I thought I could do this with a yaml like that:

```
input_features:
- name: image_path
  type: image
  preprocessing:
      num_channels: 1
      height: 250
      width: 250
  encoder:
      type: resnet
      use_pretrained: false
      trainable: true
      model_variant: 34

output_features:
- name: class
  type: category

trainer:
  epochs: 50
  batch_size: auto
  early_stop: 10
  learning_rate: 0.004
  learning_rate_scheduler:
      reduce_on_plateau: 4
      reduce_on_plateau_patience: 4
      reduce_plateau_rate: 0.5
```

This should, so I thought, train a ResNet34 from scratch. However, the training does not start since there is a conflict between output and broadcast shapes. That's also the reason why I defined my image dimensions manually, which was never required before. I still get:

```
[1]<stderr>:  File "/home/ray/anaconda3/bin/ludwig", line 8, in <module>
[1]<stderr>:    sys.exit(main())
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/cli.py", line 172, in main
[1]<stderr>:    CLI()
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/cli.py", line 67, in __init__
[1]<stderr>:    getattr(self, args.command)()
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/cli.py", line 72, in train
[1]<stderr>:    train.cli(sys.argv[2:])
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/train.py", line 395, in cli
[1]<stderr>:    train_cli(**vars(args))
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/train.py", line 185, in train_cli
[1]<stderr>:    model.train(
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/api.py", line 535, in train
[1]<stderr>:    preprocessed_data = self.preprocess(
[1]<stderr>:  File "/home/ray/anaconda3/lib/python3.8/site-packages/ludwig/api.py", line 1428, in preprocess
[1]<stderr>:    raise RuntimeError(f"Caught exception during model preprocessing: {str(e)}") from e
[1]<stderr>:RuntimeError: Caught exception during model preprocessing: output with shape [1, 224, 224] doesn't match the broadcast shape [3, 224, 224]
```

What puzzles me is the shape conflict that has nothing to do with the actual dimensions of my data. I get the same error with the `resnext` encoder. Do I have to resize my data accordingly?

What works however (because it worked before) is using the old `_resnet_legacy` encoder with a yaml below. I assumed naively I can use the TorchVision encoders the same way. but I might be wrong there.

```
input_features:
- name: image_path
  type: image
  preprocessing:
      num_channels: 1
      height: 250
      width: 250
  encoder:
    type: _resnet_legacy
    dropout: 0.0
    output_size: 128
    activation: relu
    norm: null
    first_pool_kernel_size: null
    first_pool_stride: null
    use_bias: true
    bias_initializer: zeros
    weights_initializer: xavier_uniform
    norm_params: null
    num_fc_layers: 1
    fc_layers: null
    resnet_size: 50
    num_channels: null
    out_channels: 32
    kernel_size: 3
    conv_stride: 1
    batch_norm_momentum: 0.9
    batch_norm_epsilon: 0.001

output_features:
- name: class
  type: category

trainer:
  epochs: 50
  batch_size: auto
  early_stop: 10
  learning_rate: 0.004
  learning_rate_scheduler:
      reduce_on_plateau: 4
      reduce_on_plateau_patience: 4
      reduce_plateau_rate: 0.5
```

Any feedback for clarification is much appreciated. I am using the ludwig-ray-gpu:master docker container (Off topic: this container does not seem to contain `mlflow` and requires a manual installation to avoid an early training crash).

Hints from issue discussion:
Hey @tboo, thanks for reporting this. The issue here is that the pretrained TorchVision models all assume 3 channel input images. I consider this a bug on our side, since it should work without you needing to do anything in particular to your data or adjust preprocessing.

At the moment I'm thinking the fix on our side should be to automatically replicate the single channel as 3 channels during preprocessing to feed into the TorchVision model (similar to how we crop to 224x224 during preprocessing to fit the input expectations). That way it would just work transparently without needing to modify data or preprocessing.

cc @jimthompson5802 

Regarding mlflow, you are correct that we don't include it my default, since it's considered a "contrib" module (so we don't list it in any of our requirements). However, this could be a reasonable thing to do since a lot of users are using mlflow. I can make a note to fix that.

Commit message: add descriptive docstring

Patch to review:
diff --git a/ludwig/features/image_feature.py b/ludwig/features/image_feature.py
index 50388fda3aa..48cb4f30ce2 100644
--- a/ludwig/features/image_feature.py
+++ b/ludwig/features/image_feature.py
@@ -47,6 +47,7 @@
     WIDTH,
 )
 from ludwig.data.cache.types import wrap
+from ludwig.encoders.image.torchvision import TVModelVariant
 from ludwig.features.base_feature import BaseFeatureMixin, InputFeature
 from ludwig.schema.features.augmentation.base import BaseAugmentationConfig
 from ludwig.schema.features.augmentation.image import (
@@ -71,6 +72,7 @@
     read_image_from_bytes_obj,
     read_image_from_path,
     resize_image,
+    ResizeChannels,
     torchvision_model_registry,
 )
 from ludwig.utils.misc_utils import set_default_value
@@ -241,16 +243,56 @@ def _renormalize_image(self, images):
             return images.type(torch.float32).div(255.0)
 
 
+def _get_torchvision_transform(torchvision_parameters: TVModelVariant) -> torch.nn.Module:
+    """Returns a torchvision transform that is compatible with the model variant.
+
+    Note that the raw torchvision transform is not returned. Instead, a Sequential module that includes
+    image resizing is returned. This is because the raw torchvision transform assumes that the input image has
+    three channels, which is not always the case with images input into Ludwig.
+
+    Args:
+        torchvision_parameters: The parameters for the torchvision model variant.
+    Returns:
+        (torchvision_transform, transform_metadata): A torchvision transform and the metadata for the transform.
+    """
+    torchvision_transform_raw = torchvision_parameters.model_weights.DEFAULT.transforms()
+    torchvision_transform = torch.nn.Sequential(
+        ResizeChannels(num_channels=3),
+        torchvision_transform_raw,
+    )
+    transform_metadata = {
+        "height": torchvision_transform_raw.crop_size[0],
+        "width": torchvision_transform_raw.crop_size[0],
+        "num_channels": len(torchvision_transform_raw.mean),
+    }
+    return (torchvision_transform, transform_metadata)
+
+
+def _get_torchvision_parameters(model_type: str, model_variant: str) -> TVModelVariant:
+    return torchvision_model_registry.get(model_type).get(model_variant)
+
+
 class _ImagePreprocessing(torch.nn.Module):
     """Torchscript-enabled version of preprocessing done by ImageFeatureMixin.add_feature_data."""
 
-    def __init__(self, metadata: TrainingSetMetadataDict, tv_transforms: Optional[torch.nn.Module] = None):
+    def __init__(
+        self,
+        metadata: TrainingSetMetadataDict,
+        torchvision_transform: Optional[torch.nn.Module] = None,
+        transform_metadata: Optional[Dict[str, Any]] = None,
+    ):
         super().__init__()
-        self.height = metadata["preprocessing"]["height"]
-        self.width = metadata["preprocessing"]["width"]
-        self.num_channels = metadata["preprocessing"]["num_channels"]
+
         self.resize_method = metadata["preprocessing"]["resize_method"]
-        self.tv_transforms = tv_transforms
+        self.torchvision_transform = torchvision_transform
+        if transform_metadata is not None:
+            self.height = transform_metadata["height"]
+            self.width = transform_metadata["width"]
+            self.num_channels = transform_metadata["num_channels"]
+        else:
+            self.height = metadata["preprocessing"]["height"]
+            self.width = metadata["preprocessing"]["width"]
+            self.num_channels = metadata["preprocessing"]["num_channels"]
 
     def forward(self, v: TorchscriptPreprocessingInput) -> torch.Tensor:
         """Takes a list of images and adjusts the size and number of channels as specified in the metadata.
@@ -262,14 +304,14 @@ def forward(self, v: TorchscriptPreprocessingInput) -> torch.Tensor:
             if not torch.jit.isinstance(v, torch.Tensor):
                 raise ValueError(f"Unsupported input: {v}")
 
-        if self.tv_transforms is not None:
+        if self.torchvision_transform is not None:
             # perform pre-processing for torchvision pretrained model encoders
             if torch.jit.isinstance(v, List[torch.Tensor]):
-                imgs = [self.tv_transforms(img) for img in v]
+                imgs = [self.torchvision_transform(img) for img in v]
             else:
                 # convert batch of image tensors to a list and then run torchvision pretrained
                 # model transforms on each image
-                imgs = [self.tv_transforms(img) for img in torch.unbind(v)]
+                imgs = [self.torchvision_transform(img) for img in torch.unbind(v)]
 
             # collect the list of images into a batch
             imgs_stacked = torch.stack(imgs)
@@ -701,16 +743,23 @@ def add_feature_data(
         model_type = feature_config[ENCODER].get("type", None)
         model_variant = feature_config[ENCODER].get("model_variant")
         if model_variant:
-            torchvision_parameters = torchvision_model_registry.get(model_type).get(model_variant)
+            torchvision_parameters = _get_torchvision_parameters(model_type, model_variant)
         else:
             torchvision_parameters = None
 
         if torchvision_parameters:
+            logger.warning(
+                f"Using the transforms specified for the torchvision model {model_type} {model_variant}"
+                f"This includes setting the number of channels is 3 and resizing the image to the needs of the model."
+            )
+
+            torchvision_transform, transform_metadata = _get_torchvision_transform(torchvision_parameters)
+
             # torchvision_parameters is not None
             # perform torchvision model transformations
             read_image_if_bytes_obj_and_resize = partial(
                 ImageFeatureMixin._read_image_with_pretrained_transform,
-                transform_fn=torchvision_parameters.model_weights.DEFAULT.transforms(),
+                transform_fn=torchvision_transform,
             )
             average_file_size = None
 
@@ -724,8 +773,9 @@ def add_feature_data(
             preprocessing_parameters["torchvision_model_variant"] = model_variant
 
             # get required setup parameters for in_memory = False processing
-            width = height = read_image_if_bytes_obj_and_resize.keywords["transform_fn"].crop_size[0]
-            num_channels = len(read_image_if_bytes_obj_and_resize.keywords["transform_fn"].mean)
+            height = transform_metadata["height"]
+            width = transform_metadata["width"]
+            num_channels = transform_metadata["num_channels"]
         else:
             # torchvision_parameters is None
             # perform Ludwig specified transformations
@@ -885,16 +935,18 @@ def create_preproc_module(metadata: Dict[str, Any]) -> torch.nn.Module:
         model_type = metadata["preprocessing"].get("torchvision_model_type")
         model_variant = metadata["preprocessing"].get("torchvision_model_variant")
         if model_variant:
-            torchvision_parameters = torchvision_model_registry.get(model_type).get(model_variant)
+            torchvision_parameters = _get_torchvision_parameters(model_type, model_variant)
         else:
             torchvision_parameters = None
 
         if torchvision_parameters:
-            tv_transforms = torchvision_parameters.model_weights.DEFAULT.transforms()
+            torchvision_transform, transform_metadata = _get_torchvision_transform(torchvision_parameters)
         else:
-            tv_transforms = None
+            torchvision_transform = None
 
-        return _ImagePreprocessing(metadata, tv_transforms=tv_transforms)
+        return _ImagePreprocessing(
+            metadata, torchvision_transform=torchvision_transform, transform_metadata=transform_metadata
+        )
 
     def get_augmentation_pipeline(self):
         return self.augmentation_pipeline
diff --git a/ludwig/utils/image_utils.py b/ludwig/utils/image_utils.py
index b0126b7d01d..613eb09a1a8 100644
--- a/ludwig/utils/image_utils.py
+++ b/ludwig/utils/image_utils.py
@@ -51,6 +51,30 @@ class TVModelVariant:
 IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".tiff", ".bmp", ".gif")
 
 
+@DeveloperAPI
+class ResizeChannels(torch.nn.Module):
+    def __init__(self, num_channels: int):
+        super().__init__()
+        self.num_channels = num_channels
+
+    def forward(self, imgs: torch.Tensor):
+        original_imgs_shape = imgs.shape
+        if len(original_imgs_shape) == 3:  # if shape is (C, H, W), add batch dimension
+            imgs = imgs.unsqueeze(0)
+
+        channels = imgs.shape[1]
+        if channels > self.num_channels:
+            # take the first `self.num_channels` channels
+            imgs = imgs[:, : self.num_channels, :, :]
+        elif channels < self.num_channels:
+            # repeat and use the first `self.num_channels` channels
+            imgs = imgs.repeat(1, (self.num_channels // channels) + 1, 1, 1)[:, : self.num_channels, :, :]
+
+        if len(original_imgs_shape) == 3:  # if shape was (C, H, W), remove batch dimension
+            return imgs[0]
+        return imgs
+
+
 @DeveloperAPI
 def get_gray_default_image(num_channels: int, height: int, width: int) -> np.ndarray:
     return np.full((num_channels, height, width), 128, dtype=np.float32)
diff --git a/tests/integration_tests/test_preprocessing.py b/tests/integration_tests/test_preprocessing.py
index 689661cda00..9beb2ec667f 100644
--- a/tests/integration_tests/test_preprocessing.py
+++ b/tests/integration_tests/test_preprocessing.py
@@ -593,6 +593,33 @@ def test_vit_encoder_different_dimension_image(tmpdir, csv_filename, use_pretrai
     model.train(dataset=data_csv)
 
 
+def test_image_encoder_torchvision_different_num_channels(tmpdir, csv_filename):
+    input_features = [
+        image_feature(
+            os.path.join(tmpdir, "generated_output"),
+            preprocessing={"in_memory": True, "height": 224, "width": 206, "num_channels": 1},
+            encoder={"type": "efficientnet"},
+        )
+    ]
+    output_features = [category_feature(decoder={"vocab_size": 5}, reduce_input="sum")]
+
+    data_csv = generate_data(
+        input_features, output_features, os.path.join(tmpdir, csv_filename), num_examples=NUM_EXAMPLES
+    )
+
+    config = {
+        "input_features": input_features,
+        "output_features": output_features,
+        "trainer": {"train_steps": 1},
+    }
+
+    model = LudwigModel(config)
+
+    # Failure happens post preprocessing but before training during the ECD model creation phase
+    # so make sure the model can be created properly and training can proceed
+    model.train(dataset=data_csv)
+
+
 @pytest.mark.parametrize(
     "df_engine",
     [
diff --git a/tests/ludwig/utils/test_image_utils.py b/tests/ludwig/utils/test_image_utils.py
index 6a4a1d647d3..32254cb0543 100644
--- a/tests/ludwig/utils/test_image_utils.py
+++ b/tests/ludwig/utils/test_image_utils.py
@@ -17,7 +17,15 @@
 import pytest
 import torch
 
-from ludwig.utils.image_utils import crop, crop_or_pad, grayscale, num_channels_in_image, pad, resize_image
+from ludwig.utils.image_utils import (
+    crop,
+    crop_or_pad,
+    grayscale,
+    num_channels_in_image,
+    pad,
+    resize_image,
+    ResizeChannels,
+)
 
 
 @pytest.mark.parametrize("pad_fn", [pad, torch.jit.script(pad)])
@@ -223,3 +231,19 @@ def test_num_channels_in_image():
     with pytest.raises(ValueError):
         num_channels_in_image(torch.rand(5))
         num_channels_in_image(None)
+
+
+@pytest.mark.parametrize("image_shape", [(1, 10, 10), (3, 10, 10), (5, 10, 10)])
+@pytest.mark.parametrize("num_channels_expected", [1, 2, 3, 4])
+def test_ResizeChannels_module(image_shape, num_channels_expected):
+    image = torch.randint(0, 1, image_shape)
+    fn = ResizeChannels(num_channels_expected)
+    assert fn(image).shape == tuple([num_channels_expected] + list(image_shape[1:]))
+
+
+@pytest.mark.parametrize("image_shape", [(2, 1, 10, 10), (2, 3, 10, 10), (2, 5, 10, 10)])
+@pytest.mark.parametrize("num_channels_expected", [1, 2, 3, 4])
+def test_ResizeChannels_module_with_batch_dim(image_shape, num_channels_expected):
+    image = torch.randint(0, 1, image_shape)
+    fn = ResizeChannels(num_channels_expected)
+    assert fn(image).shape == tuple([image_shape[0], num_channels_expected] + list(image_shape[2:]))


Code review comments:
diff_hunk: @@ -701,16 +743,23 @@ def add_feature_data(
         model_type = feature_config[ENCODER].get("type", None)
         model_variant = feature_config[ENCODER].get("model_variant")
         if model_variant:
-            torchvision_parameters = torchvision_model_registry.get(model_type).get(model_variant)
+            torchvision_parameters = _get_torchvision_parameters(model_type, model_variant)
         else:
             torchvision_parameters = None
 
         if torchvision_parameters:
+            logger.warning(
+                f"Using the transforms specified for the torchvision model {model_type} {model_variant}"
path: ludwig/features/image_feature.py
text: @user1:
Nit: add whitespace at the end to prevent output like `model model_type model_variantThis include`