# -------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License (MIT). See LICENSE in the repo root for license information.
# -------------------------------------------------------------------------------------------
from __future__ import annotations
from contextlib import contextmanager
from typing import Any, Generator, Optional, Sequence, Tuple, Union
import torch
import torch.nn as nn
from health_multimodal.common.device import get_module_device
from timm.models.layers import trunc_normal_
from .resnet import resnet18, resnet50
from .transformer import VisionTransformerPooler
from .types import ImageEncoderType
DEFAULT_DILATION_VALUES_FOR_RESNET = (False, False, True)
ImageEncoderOutputType = Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]
class ImageEncoder(nn.Module):
"""Image encoder trunk module for the ``ImageModel`` class.
:param img_encoder_type : Type of image encoder model to use, either ``"resnet18_multi_image"`` or
``"resnet50_multi_image"``.
"""
def __init__(self, img_encoder_type: str):
super().__init__()
self.img_encoder_type = img_encoder_type
self.encoder = self._create_encoder()
def _create_encoder(self, **kwargs: Any) -> nn.Module:
if self.img_encoder_type in [ImageEncoderType.RESNET18, ImageEncoderType.RESNET18_MULTI_IMAGE]:
encoder_class = resnet18
elif self.img_encoder_type in [ImageEncoderType.RESNET50, ImageEncoderType.RESNET50_MULTI_IMAGE]:
encoder_class = resnet50
else:
supported = ImageEncoderType.get_members(multi_image_encoders_only=False)
raise NotImplementedError(f"Image encoder type \"{self.img_encoder_type}\" must be in {supported}")
encoder = encoder_class(pretrained=True, **kwargs)
return encoder
def forward(self,
current_image: torch.Tensor,
return_patch_embeddings: bool = False) -> ImageEncoderOutputType:
"""Get image global and patch embeddings"""
patch_emb = self.encoder(current_image)
avg_pooled_emb = torch.flatten(torch.nn.functional.adaptive_avg_pool2d(patch_emb, (1, 1)), 1)
if return_patch_embeddings:
return patch_emb, avg_pooled_emb
return avg_pooled_emb
def reload_encoder_with_dilation(self, replace_stride_with_dilation: Optional[Sequence[bool]] = None) -> None:
"""Workaround for enabling dilated convolutions after model initialization.
:param replace_stride_with_dilation: Replace the 2x2 standard convolution stride with a dilated convolution
in each layer in the last three blocks of ResNet architecture.
"""
if self.img_encoder_type == ImageEncoderType.RESNET18:
# resnet18 uses BasicBlock implementation, which does not support dilated convolutions.
raise NotImplementedError("resnet18 does not support dilated convolutions")
if replace_stride_with_dilation is None:
replace_stride_with_dilation = DEFAULT_DILATION_VALUES_FOR_RESNET
device = next(self.encoder.parameters()).device
new_encoder = self._create_encoder(replace_stride_with_dilation=replace_stride_with_dilation).to(device)
if self.encoder.training:
new_encoder.train()
else:
new_encoder.eval()
new_encoder.load_state_dict(self.encoder.state_dict())
self.encoder = new_encoder
class MultiImageEncoder(ImageEncoder):
"""Multi-image encoder trunk module for the ``ImageModel`` class.
It can be used to encode multiple images into combined latent representation.
Currently it only supports two input images but can be extended to support more in future.
:param img_encoder_type: Type of image encoder model to use: either ``"resnet18"`` or ``"resnet50"``.
"""
def __init__(self, img_encoder_type: str):
super().__init__(img_encoder_type)
output_dim = 256 # The aggregate feature dim of the encoder is `2 * output_dim` i.e. [f_static, f_diff]
grid_shape = (14, 14) # Spatial dimensions of patch grid.
backbone_output_feature_dim = get_encoder_output_dim(self.encoder, device=get_module_device(self))
self.backbone_to_vit = nn.Conv2d(in_channels=backbone_output_feature_dim, out_channels=output_dim,
kernel_size=1, stride=1, padding=0, bias=False)
self.vit_pooler = VisionTransformerPooler(input_dim=output_dim, grid_shape=grid_shape)
# Missing image embedding
self.missing_previous_emb = nn.Parameter(torch.zeros(1, output_dim, 1, 1))
trunc_normal_(self.missing_previous_emb, std=.02)
def forward(self, # type: ignore[override]
current_image: torch.Tensor,
previous_image: Optional[torch.Tensor] = None,
return_patch_embeddings: bool = False) -> ImageEncoderOutputType:
batch_size = current_image.shape[0]
if previous_image is not None:
assert current_image.shape == previous_image.shape
x = torch.cat([current_image, previous_image], dim=0)
x = super().forward(x, return_patch_embeddings=True)[0]
x = self.backbone_to_vit(x)
patch_x, patch_x_previous = x[:batch_size], x[batch_size:]
diff_x = self.vit_pooler(current_image=patch_x, previous_image=patch_x_previous)
else:
x = super().forward(current_image, return_patch_embeddings=True)[0]
patch_x = self.backbone_to_vit(x)
B, _, W, H = patch_x.shape
diff_x = self.missing_previous_emb.repeat(B, 1, W, H)
patch_fused = torch.cat([patch_x, diff_x], dim=1)
avg_pooled_emb = torch.flatten(torch.nn.functional.adaptive_avg_pool2d(patch_fused, (1, 1)), 1)
if return_patch_embeddings:
return patch_fused, avg_pooled_emb
return avg_pooled_emb
def reload_encoder_with_dilation(self, replace_stride_with_dilation: Optional[Sequence[bool]] = None) -> None:
raise NotImplementedError
@torch.no_grad()
def get_encoder_output_dim(module: torch.nn.Module, device: torch.device) -> int:
"""Calculate the output dimension of an encoder by making a single forward pass.
:param module: Encoder module.
:param device: Compute device to use.
"""
# Target device
assert isinstance(device, torch.device)
x = torch.rand((1, 3, 448, 448)).to(device)
# Extract the number of output feature dimensions
with restore_training_mode(module):
module.eval()
representations = module(x)
return representations.shape[1]
@contextmanager
def restore_training_mode(module: nn.Module) -> Generator[None, None, None]:
"""Restore the training mode of a module after some operation.
:param module: PyTorch module.
"""
training_mode = module.training
yield
module.train(mode=training_mode)
def get_encoder_from_type(img_encoder_type: str) -> ImageEncoder:
"""Returns the encoder class for the given encoder type.
:param img_encoder_type: Encoder type. {RESNET18, RESNET50, RESNET18_MULTI_IMAGE, RESNET50_MULTI_IMAGE}
"""
if img_encoder_type in ImageEncoderType.get_members(multi_image_encoders_only=True):
return MultiImageEncoder(img_encoder_type=img_encoder_type)
else:
return ImageEncoder(img_encoder_type=img_encoder_type)
