77 lines
3.3 KiB
Python
77 lines
3.3 KiB
Python
"""ResNet34 backbone + 3-scale FPN-style fusion.
|
|
|
|
Reproduces the paper's genuinely distinctive idea (multi-scale down/up-sample fusion
|
|
to help extract thin, elongated lane structure) with a concretely-defined FPN instead
|
|
of the paper's ambiguous "32-group ResNet32/ResNeXt50" description.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torchvision
|
|
|
|
|
|
class ConvBNMish(nn.Module):
|
|
def __init__(self, in_ch: int, out_ch: int, kernel_size: int = 3, padding: int = 1):
|
|
super().__init__()
|
|
self.conv = nn.Conv2d(in_ch, out_ch, kernel_size, padding=padding, bias=False)
|
|
self.bn = nn.BatchNorm2d(out_ch)
|
|
self.act = nn.Mish(inplace=True)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return self.act(self.bn(self.conv(x)))
|
|
|
|
|
|
class Backbone(nn.Module):
|
|
"""Fuses stride 8/16/32 features (paper's 3-scale idea), then downsamples the
|
|
fused map to stride 32 before returning it -- full O(N^2) self-attention over a
|
|
stride-8 map (e.g. 45x80=3600 tokens) is prohibitively expensive/OOMs on CPU;
|
|
the paper's own tiny FLOPs count (0.425 GMACs) implies the same heavy
|
|
downsampling before attention, just left unstated. Stride 32 on a 360x640 input
|
|
gives an 11x20=220-token sequence, ~260x cheaper attention.
|
|
"""
|
|
|
|
def __init__(self, name: str = "resnet34", pretrained: bool = True, out_channels: int = 128):
|
|
super().__init__()
|
|
weights = torchvision.models.ResNet34_Weights.IMAGENET1K_V1 if pretrained else None
|
|
net = torchvision.models.resnet34(weights=weights)
|
|
|
|
self.stem = nn.Sequential(net.conv1, net.bn1, net.relu, net.maxpool)
|
|
self.layer1 = net.layer1 # stride 4
|
|
self.layer2 = net.layer2 # stride 8, C=128
|
|
self.layer3 = net.layer3 # stride 16, C=256
|
|
self.layer4 = net.layer4 # stride 32, C=512
|
|
|
|
# Shared extractor delta: project each scale to a common channel dim.
|
|
self.reduce_c3 = ConvBNMish(128, out_channels, kernel_size=1, padding=0)
|
|
self.reduce_c4 = ConvBNMish(256, out_channels, kernel_size=1, padding=0)
|
|
self.reduce_c5 = ConvBNMish(512, out_channels, kernel_size=1, padding=0)
|
|
|
|
self.fuse_c4 = ConvBNMish(out_channels, out_channels)
|
|
self.fuse_c3 = ConvBNMish(out_channels, out_channels)
|
|
|
|
# stride 8 -> stride 32 for the transformer input (see class docstring).
|
|
self.downsample = nn.Sequential(
|
|
nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
|
|
nn.BatchNorm2d(out_channels), nn.Mish(inplace=True),
|
|
nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
|
|
nn.BatchNorm2d(out_channels), nn.Mish(inplace=True),
|
|
)
|
|
|
|
self.out_channels = out_channels
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = self.stem(x)
|
|
x = self.layer1(x)
|
|
c3 = self.layer2(x) # stride 8
|
|
c4 = self.layer3(c3) # stride 16
|
|
c5 = self.layer4(c4) # stride 32
|
|
|
|
p5 = self.reduce_c5(c5)
|
|
p4 = self.reduce_c4(c4) + nn.functional.interpolate(p5, size=c4.shape[-2:], mode="nearest")
|
|
p4 = self.fuse_c4(p4)
|
|
p3 = self.reduce_c3(c3) + nn.functional.interpolate(p4, size=c3.shape[-2:], mode="nearest")
|
|
p3 = self.fuse_c3(p3)
|
|
|
|
return self.downsample(p3) # (B, out_channels, H/32, W/32)
|