分享几种resnet框架的backbone代码（pytorch）
import torchvision
import torch.nn as nn
__all__ = ["ResNet18", "ResNet50","resnext50_32x4d"]
class ResNet18(nn.Module):
output_size = 512
def __init__(self, pretrained=True):
super(ResNet18, self).__init__()
pretrained = torchvision.models.resnet18(pretrained=pretrained)
for module_name in [
"conv1",
"bn1",
"relu",
"maxpool",
"layer1",
"layer2",
"layer3",
"layer4",
"avgpool",
]:
self.add_module(module_name, getattr(pretrained, module_name))
def forward(self, x, get_ha=False):
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
b1 = self.layer1(x)
b2 = self.layer2(b1)
b3 = self.layer3(b2)
b4 = self.layer4(b3)
pool = self.avgpool(b4)
if get_ha:
return b1, b2, b3, b4, pool
return pool
class ResNet34(nn.Module):
output_size = 512
def __init__(self, pretrained=True):
super(ResNet34, self).__init__()
pretrained = torchvision.models.resnet34(pretrained=pretrained)
for module_name in [
"conv1",
"bn1",
"relu",
"maxpool",
"layer1",
"layer2",
"layer3",
"layer4",
"avgpool",
]:
self.add_module(module_name, getattr(pretrained, module_name))
def forward(self, x, get_ha=False):
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
b1 = self.layer1(x)
b2 = self.layer2(b1)
b3 = self.layer3(b2)
b4 = self.layer4(b3)
pool = self.avgpool(b4)
if get_ha:
return b1, b2, b3, b4, pool
return pool
class ResNet50(nn.Module):
output_size = 2048
def __init__(self, pretrained=True):
super(ResNet50, self).__init__()
pretrained = torchvision.models.resnet50(pretrained=pretrained)
for module_name in [
"conv1",
"bn1",
"relu",
"maxpool",
"layer1",
"layer2",
"layer3",
"layer4",
"avgpool",
]:
self.add_module(module_name, getattr(pretrained, module_name))
def forward(self, x, get_ha=False):
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
b1 = self.layer1(x)
b2 = self.layer2(b1)
b3 = self.layer3(b2)
b4 = self.layer4(b3)
pool = self.avgpool(b4)
if get_ha:
return b1, b2, b3, b4, pool
return pool
class resnext50_32x4d(nn.Module):
output_size = 2048
def __init__(self, pretrained=True):
super(resnext50_32x4d, self).__init__()
pretrained = torchvision.models.resnext50_32x4d(pretrained=pretrained)
for module_name in [
"conv1",
"bn1",
"relu",
"maxpool",
"layer1",
"layer2",
"layer3",
"layer4",
"avgpool",
]:
self.add_module(module_name, getattr(pretrained, module_name))
def forward(self, x, get_ha=False):
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
b1 = self.layer1(x)
b2 = self.layer2(b1)
b3 = self.layer3(b2)
b4 = self.layer4(b3)
pool = self.avgpool(b4)
if get_ha:
return b1, b2, b3, b4, pool
return pool
class resnet152(nn.Module):
output_size = 2048
def __init__(self, pretrained=True):
super(resnet152, self).__init__()
pretrained = torchvision.models.resnet152(pretrained=pretrained)
for module_name in [
"conv1",
"bn1",
"relu",
"maxpool",
"layer1",
"layer2",
"layer3",
"layer4",
"avgpool",
]:
self.add_module(module_name, getattr(pretrained, module_name))
def forward(self, x, get_ha=False):
x = self.maxpool(self.relu(self.bn1(self.conv1(x))))
b1 = self.layer1(x)
b2 = self.layer2(b1)
b3 = self.layer3(b2)
b4 = self.layer4(b3)
pool = self.avgpool(b4)
if get_ha:
return b1, b2, b3, b4, pool
return pool可以用于再深度学习训练过程中对网络框架集成或者更改。