<filename>models/backbones/resnet.py
import torch
import torch.nn as nn
from base.base_model import BaseBackbone


def conv3x3(in_planes, out_planes, stride=1, dilation=1):
	return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride, padding=dilation, dilation=dilation, bias=False)

def conv1x1(in_planes, out_planes, stride=1):
	return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)

class BasicBlock(nn.Module):
	expansion = 1

	def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
		super(BasicBlock, self).__init__()
		self.conv1 = conv3x3(inplanes, planes, stride, dilation=dilation)
		self.bn1 = nn.BatchNorm2d(planes)
		self.relu = nn.ReLU(inplace=True)

		self.conv2 = conv3x3(planes, planes, dilation=dilation)
		self.bn2 = nn.BatchNorm2d(planes)

		self.downsample = downsample
		self.stride = stride

	def forward(self, x):
		residual = x

		out = self.conv1(x)
		out = self.bn1(out)
		out = self.relu(out)

		out = self.conv2(out)
		out = self.bn2(out)

		if self.downsample is not None:
			residual = self.downsample(x)

		out += residual
		out = self.relu(out)

		return out

class Bottleneck(nn.Module):
	expansion = 4

	def __init__(self, inplanes, planes, stride=1, downsample=None, dilation=1):
		super(Bottleneck, self).__init__()
		self.conv1 = conv1x1(inplanes, planes)
		self.bn1 = nn.BatchNorm2d(planes)

		self.conv2 = conv3x3(planes, planes, stride, dilation=dilation)
		self.bn2 = nn.BatchNorm2d(planes)

		self.conv3 = conv1x1(planes, planes * self.expansion)
		self.bn3 = nn.BatchNorm2d(planes * self.expansion)
		
		self.relu = nn.ReLU(inplace=True)
		self.downsample = downsample
		self.stride = stride

	def forward(self, x):
		residual = x

		out = self.conv1(x)
		out = self.bn1(out)
		out = self.relu(out)

		out = self.conv2(out)
		out = self.bn2(out)
		out = self.relu(out)

		out = self.conv3(out)
		out = self.bn3(out)

		if self.downsample is not None:
			residual = self.downsample(x)

		out += residual
		out = self.relu(out)

		return out

class ResNet(BaseBackbone):
	basic_inplanes = 64

	def __init__(self, block, layers, output_stride=32, num_classes=1000):
		super(ResNet, self).__init__()
		self.inplanes = self.basic_inplanes
		self.output_stride = output_stride
		self.num_classes = num_classes

		if output_stride==8:
			strides   = [1, 2, 1, 1]
			dilations = [1, 1, 2, 4]
		elif output_stride==16:
			strides   = [1, 2, 2, 1]
			dilations = [1, 1, 1, 2]
		elif output_stride==32:
			strides   = [1, 2, 2, 2]
			dilations = [1, 1, 1, 1]
		else:
			raise NotImplementedError

		self.conv1 = nn.Conv2d(3, self.basic_inplanes, kernel_size=7, stride=2, padding=3, bias=False)
		self.bn1 = nn.BatchNorm2d(self.basic_inplanes)
		self.relu = nn.ReLU(inplace=True)
		self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)

		self.layer1 = self._make_layer(block, 1*self.basic_inplanes, num_layers=layers[0], stride=strides[0], dilation=dilations[0])
		self.layer2 = self._make_layer(block, 2*self.basic_inplanes, num_layers=layers[1], stride=strides[1], dilation=dilations[1])
		self.layer3 = self._make_layer(block, 4*self.basic_inplanes, num_layers=layers[2], stride=strides[2], dilation=dilations[2])
		self.layer4 = self._make_layer(block, 8*self.basic_inplanes, num_layers=layers[3], stride=strides[3], dilation=dilations[3])

		if self.num_classes is not None:
			self.fc = nn.Linear(8*self.basic_inplanes * block.expansion, num_classes)

		self.init_weights()

	def forward(self, x):
		# Stage1
		x = self.conv1(x)
		x = self.bn1(x)
		x = self.relu(x)
		# Stage2
		x = self.maxpool(x)
		x = self.layer1(x)
		# Stage3
		x = self.layer2(x)
		# Stage4
		x = self.layer3(x)
		# Stage5
		x = self.layer4(x)
		# Classification
		if self.num_classes is not None:
			x = x.mean(dim=(2,3))
			x = self.fc(x)
		# Output
		return x

	def _make_layer(self, block, planes, num_layers, stride=1, dilation=1, grids=None):
		# Downsampler
		downsample = None
		if (stride != 1) or (self.inplanes != planes * block.expansion):
			downsample = nn.Sequential(
				conv1x1(self.inplanes, planes * block.expansion, stride),
				nn.BatchNorm2d(planes * block.expansion))
		# Multi-grids
		if dilation!=1:
			dilations = [dilation*(2**layer_idx) for layer_idx in range(num_layers)]
		else:
			dilations = num_layers*[dilation]
		# Construct layers
		layers = []
		layers.append(block(self.inplanes, planes, stride, downsample, dilations[0]))
		self.inplanes = planes * block.expansion
		for i in range(1, num_layers):
			layers.append(block(self.inplanes, planes, dilation=dilations[i]))
		return nn.Sequential(*layers)

def resnet18(pretrained=None, **kwargs):
	model = ResNet(BasicBlock, [2, 2, 2, 2], **kwargs)
	if pretrained is not None:
		model._load_pretrained_model(pretrained)
	return model

def resnet34(pretrained=None, **kwargs):
	model = ResNet(BasicBlock, [3, 4, 6, 3], **kwargs)
	if pretrained is not None:
		model._load_pretrained_model(pretrained)
	return model

def resnet50(pretrained=None, **kwargs):
	model = ResNet(Bottleneck, [3, 4, 6, 3], **kwargs)
	if pretrained is not None:
		model._load_pretrained_model(pretrained)
	return model

def resnet101(pretrained=None, **kwargs):
	model = ResNet(Bottleneck, [3, 4, 23, 3], **kwargs)
	if pretrained is not None:
		model._load_pretrained_model(pretrained)
	return model

def resnet152(pretrained=None, **kwargs):
	model = ResNet(Bottleneck, [3, 8, 36, 3], **kwargs)
	if pretrained is not None:
		model._load_pretrained_model(pretrained)
	return model

def get_resnet(num_layers, **kwargs):
	if num_layers==18:
		return resnet18(**kwargs)
	elif num_layers==34:
		return resnet34(**kwargs)
	elif num_layers==50:
		return resnet50(**kwargs)
	elif num_layers==101:
		return resnet101(**kwargs)
	elif num_layers==152:
		return resnet152(**kwargs)
	else:
		raise NotImplementedError