Add ResNet-50 as a backbone option

a2999cc about 1 year ago

4.84 kB

"""Encoder wrappers with adjustable input channels.
Supports two backbone families:
- HuggingFace Transformers SegFormer (e.g., "mit_b2")
- TorchVision ResNet-50 (use backbone "resnet50" | "resnet-50" | "resnet_50")
Both return a list of 4 multi-scale feature maps [C1, C2, C3, C4] at strides
1/4, 1/8, 1/16, 1/32 respectively.
"""

from typing import List, Tuple

import torch
import torch.nn as nn
from torchvision.models import resnet50, ResNet50_Weights


class SegFormerEncoder(nn.Module):
    def __init__(
        self,
        backbone: str = "mit_b2",
        in_channels: int = 6,
        pretrained: bool = True,
    ):
        super().__init__()
        self.backbone_name = backbone
        self.in_channels = in_channels
        self.pretrained = pretrained

        self.hf = None
        self.resnet = None

        # SegFormer path
        if backbone.startswith("mit_") or backbone.startswith("segformer"):
            self.hf = _HFEncoderWrapper(in_channels, backbone, pretrained)
            self.feature_dims = self.hf.feature_dims
        # ResNet-50 path
        elif backbone in ("resnet50", "resnet-50", "resnet_50"):
            self.resnet = _ResNetEncoderWrapper(in_channels, pretrained)
            self.feature_dims = self.resnet.feature_dims
        else:
            raise ValueError(
                f"Unsupported backbone '{backbone}'. Use one of: mit_b[0-5], segformer*, resnet50."
            )

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        if self.hf is not None:
            return self.hf(x)
        if self.resnet is not None:
            return self.resnet(x)
        raise AssertionError("No encoder instantiated")


class _ResNetEncoderWrapper(nn.Module):
    def __init__(self, in_chans: int, pretrained: bool):
        super().__init__()
        # Build base ResNet-50
        if pretrained:
            self.model = resnet50(weights=ResNet50_Weights.DEFAULT)
        else:
            self.model = resnet50(weights=None)

        # Adjust input stem for arbitrary channel count
        if in_chans != 3:
            old_conv = self.model.conv1
            new_conv = nn.Conv2d(
                in_chans, old_conv.out_channels, kernel_size=old_conv.kernel_size[0],
                stride=old_conv.stride[0], padding=old_conv.padding[0], bias=False
            )
            with torch.no_grad():
                if pretrained and old_conv.weight.shape[1] == 3:
                    w = old_conv.weight  # [64, 3, 7, 7]
                    if in_chans > 3:
                        w_mean = w.mean(dim=1, keepdim=True)
                        new_w = w_mean.repeat(1, in_chans, 1, 1)
                    else:
                        new_w = w[:, :in_chans, :, :]
                    new_conv.weight.copy_(new_w)
                else:
                    nn.init.kaiming_normal_(new_conv.weight, mode="fan_out", nonlinearity="relu")
            self.model.conv1 = new_conv

        self.feature_dims = [256, 512, 1024, 2048]

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        # Stem
        x = self.model.conv1(x)
        x = self.model.bn1(x)
        x = self.model.relu(x)
        x = self.model.maxpool(x)  # 1/4

        # Stages
        c1 = self.model.layer1(x)  # 1/4, 256
        c2 = self.model.layer2(c1)  # 1/8, 512
        c3 = self.model.layer3(c2)  # 1/16, 1024
        c4 = self.model.layer4(c3)  # 1/32, 2048
        return [c1, c2, c3, c4]


class _HFEncoderWrapper(nn.Module):
    def __init__(self, in_chans: int, backbone: str, pretrained: bool):
        super().__init__()
        # Lazy import to avoid hard dependency during tests if not used
        from transformers import SegformerModel, SegformerConfig

        name_map = {
            "mit_b0": "nvidia/mit-b0",
            "mit_b1": "nvidia/mit-b1",
            "mit_b2": "nvidia/mit-b2",
            "mit_b3": "nvidia/mit-b3",
            "mit_b4": "nvidia/mit-b4",
            "mit_b5": "nvidia/mit-b5",
        }
        model_id = name_map[backbone]

        if pretrained:
            base_cfg = SegformerConfig.from_pretrained(model_id)
            base_cfg.num_channels = in_chans
            self.model = SegformerModel.from_pretrained(
                model_id, config=base_cfg, ignore_mismatched_sizes=True
            )
        else:
            cfg = SegformerConfig()  # default config (B0-like)
            cfg.num_channels = in_chans
            self.model = SegformerModel(cfg)

        # Expose channel dims per stage
        self.feature_dims = list(self.model.config.hidden_sizes)

    def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
        outputs = self.model(
            pixel_values=x, output_hidden_states=True, return_dict=True
        )
        feats = list(outputs.hidden_states)
        assert len(feats) == 4
        return feats

	"""Encoder wrappers with adjustable input channels.

	Supports two backbone families:
	- HuggingFace Transformers SegFormer (e.g., "mit_b2")
	- TorchVision ResNet-50 (use backbone "resnet50" \| "resnet-50" \| "resnet_50")

	Both return a list of 4 multi-scale feature maps [C1, C2, C3, C4] at strides
	1/4, 1/8, 1/16, 1/32 respectively.
	"""

	from typing import List, Tuple

	import torch
	import torch.nn as nn
	from torchvision.models import resnet50, ResNet50_Weights


	class SegFormerEncoder(nn.Module):
	def __init__(
	self,
	backbone: str = "mit_b2",
	in_channels: int = 6,
	pretrained: bool = True,
	):
	super().__init__()
	self.backbone_name = backbone
	self.in_channels = in_channels
	self.pretrained = pretrained

	self.hf = None
	self.resnet = None

	# SegFormer path
	if backbone.startswith("mit_") or backbone.startswith("segformer"):
	self.hf = _HFEncoderWrapper(in_channels, backbone, pretrained)
	self.feature_dims = self.hf.feature_dims
	# ResNet-50 path
	elif backbone in ("resnet50", "resnet-50", "resnet_50"):
	self.resnet = _ResNetEncoderWrapper(in_channels, pretrained)
	self.feature_dims = self.resnet.feature_dims
	else:
	raise ValueError(
	f"Unsupported backbone '{backbone}'. Use one of: mit_b[0-5], segformer*, resnet50."
	)

	def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
	if self.hf is not None:
	return self.hf(x)
	if self.resnet is not None:
	return self.resnet(x)
	raise AssertionError("No encoder instantiated")


	class _ResNetEncoderWrapper(nn.Module):
	def __init__(self, in_chans: int, pretrained: bool):
	super().__init__()
	# Build base ResNet-50
	if pretrained:
	self.model = resnet50(weights=ResNet50_Weights.DEFAULT)
	else:
	self.model = resnet50(weights=None)

	# Adjust input stem for arbitrary channel count
	if in_chans != 3:
	old_conv = self.model.conv1
	new_conv = nn.Conv2d(
	in_chans, old_conv.out_channels, kernel_size=old_conv.kernel_size[0],
	stride=old_conv.stride[0], padding=old_conv.padding[0], bias=False
	)
	with torch.no_grad():
	if pretrained and old_conv.weight.shape[1] == 3:
	w = old_conv.weight # [64, 3, 7, 7]
	if in_chans > 3:
	w_mean = w.mean(dim=1, keepdim=True)
	new_w = w_mean.repeat(1, in_chans, 1, 1)
	else:
	new_w = w[:, :in_chans, :, :]
	new_conv.weight.copy_(new_w)
	else:
	nn.init.kaiming_normal_(new_conv.weight, mode="fan_out", nonlinearity="relu")
	self.model.conv1 = new_conv

	self.feature_dims = [256, 512, 1024, 2048]

	def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
	# Stem
	x = self.model.conv1(x)
	x = self.model.bn1(x)
	x = self.model.relu(x)
	x = self.model.maxpool(x) # 1/4

	# Stages
	c1 = self.model.layer1(x) # 1/4, 256
	c2 = self.model.layer2(c1) # 1/8, 512
	c3 = self.model.layer3(c2) # 1/16, 1024
	c4 = self.model.layer4(c3) # 1/32, 2048
	return [c1, c2, c3, c4]


	class _HFEncoderWrapper(nn.Module):
	def __init__(self, in_chans: int, backbone: str, pretrained: bool):
	super().__init__()
	# Lazy import to avoid hard dependency during tests if not used
	from transformers import SegformerModel, SegformerConfig

	name_map = {
	"mit_b0": "nvidia/mit-b0",
	"mit_b1": "nvidia/mit-b1",
	"mit_b2": "nvidia/mit-b2",
	"mit_b3": "nvidia/mit-b3",
	"mit_b4": "nvidia/mit-b4",
	"mit_b5": "nvidia/mit-b5",
	}
	model_id = name_map[backbone]

	if pretrained:
	base_cfg = SegformerConfig.from_pretrained(model_id)
	base_cfg.num_channels = in_chans
	self.model = SegformerModel.from_pretrained(
	model_id, config=base_cfg, ignore_mismatched_sizes=True
	)
	else:
	cfg = SegformerConfig() # default config (B0-like)
	cfg.num_channels = in_chans
	self.model = SegformerModel(cfg)

	# Expose channel dims per stage
	self.feature_dims = list(self.model.config.hidden_sizes)

	def forward(self, x: torch.Tensor) -> List[torch.Tensor]:
	outputs = self.model(
	pixel_values=x, output_hidden_states=True, return_dict=True
	)
	feats = list(outputs.hidden_states)
	assert len(feats) == 4
	return feats