shunliwang

update

8bc3305 6 months ago

2.19 kB

# author: Xinyuzhou
# email: Xinyuzhou@sjtu.edu.cn
# date: 2025-12-24
#
# The code is for ResNet50 backbone.

import os
import logging
from typing import Union
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from metrics.registry import BACKBONE

logger = logging.getLogger(__name__)

@BACKBONE.register_module(module_name="resnet50")
class ResNet50(nn.Module):
    def __init__(self, resnet_config):
        super(ResNet50, self).__init__()
        """ Constructor
        Args:
            resnet_config: configuration file with the dict format
        """
        self.num_classes = resnet_config["num_classes"]
        inc = resnet_config["inc"]
        self.mode = resnet_config["mode"]

        # Load the pretrained ResNet50 weights from torchvision
        # New API for torchvision >= 0.13:
        #   resnet = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)
        # The legacy API is still available:
        resnet = torchvision.models.resnet50(pretrained=True)

        # If you want to support inc != 3, uncomment this and adjust accordingly
        # resnet.conv1 = nn.Conv2d(inc, 64, kernel_size=7, stride=2, padding=3, bias=False)

        # Remove the final avgpool and fc layers, keeping feature maps up to layer4
        self.resnet = torch.nn.Sequential(*list(resnet.children())[:-2])

        # The final feature dimension of ResNet50 is 2048
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(2048, self.num_classes)

        if self.mode == 'adjust_channel':
            self.adjust_channel = nn.Sequential(
                nn.Conv2d(2048, 2048, 1, 1),
                nn.BatchNorm2d(2048),
                nn.ReLU(inplace=True),
            )

    def features(self, inp):
        x = self.resnet(inp)
        if self.mode == 'adjust_channel':
            x = self.adjust_channel(x)
        return x

    def classifier(self, features):
        x = self.avgpool(features)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

    def forward(self, inp):
        x = self.features(inp)
        out = self.classifier(x)
        return out

	# author: Xinyuzhou
	# email: Xinyuzhou@sjtu.edu.cn
	# date: 2025-12-24
	#
	# The code is for ResNet50 backbone.

	import os
	import logging
	from typing import Union
	import torch
	import torchvision
	import torch.nn as nn
	import torch.nn.functional as F
	from metrics.registry import BACKBONE

	logger = logging.getLogger(__name__)

	@BACKBONE.register_module(module_name="resnet50")
	class ResNet50(nn.Module):
	def __init__(self, resnet_config):
	super(ResNet50, self).__init__()
	""" Constructor
	Args:
	resnet_config: configuration file with the dict format
	"""
	self.num_classes = resnet_config["num_classes"]
	inc = resnet_config["inc"]
	self.mode = resnet_config["mode"]

	# Load the pretrained ResNet50 weights from torchvision
	# New API for torchvision >= 0.13:
	# resnet = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.IMAGENET1K_V1)
	# The legacy API is still available:
	resnet = torchvision.models.resnet50(pretrained=True)

	# If you want to support inc != 3, uncomment this and adjust accordingly
	# resnet.conv1 = nn.Conv2d(inc, 64, kernel_size=7, stride=2, padding=3, bias=False)

	# Remove the final avgpool and fc layers, keeping feature maps up to layer4
	self.resnet = torch.nn.Sequential(*list(resnet.children())[:-2])

	# The final feature dimension of ResNet50 is 2048
	self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
	self.fc = nn.Linear(2048, self.num_classes)

	if self.mode == 'adjust_channel':
	self.adjust_channel = nn.Sequential(
	nn.Conv2d(2048, 2048, 1, 1),
	nn.BatchNorm2d(2048),
	nn.ReLU(inplace=True),
	)

	def features(self, inp):
	x = self.resnet(inp)
	if self.mode == 'adjust_channel':
	x = self.adjust_channel(x)
	return x

	def classifier(self, features):
	x = self.avgpool(features)
	x = x.view(x.size(0), -1)
	x = self.fc(x)
	return x

	def forward(self, inp):
	x = self.features(inp)
	out = self.classifier(x)
	return out