cvmodelz/cvmodelz/models/pretrained/resnet.py

import chainer

from chainer import functions as F
from chainer import links as L
from chainer.links.model.vision.resnet import BuildingBlock
from chainer.links.model.vision.resnet import _global_average_pooling_2d
from collections import OrderedDict
from functools import partial


from cvmodelz.models.meta_info import ModelInfo
from cvmodelz.models.pretrained.base import PretrainedModelMixin

class BaseResNet(PretrainedModelMixin):
	n_layers = ""

	def __init__(self, *args, **kwargs):
		super(BaseResNet, self).__init__(*args, pooling=_global_average_pooling_2d, **kwargs)
		self.meta = ModelInfo(
			name=f"ResNet{self.n_layers}",
			input_size=224,
			feature_size=2048,
			n_conv_maps=2048,

			conv_map_layer="res5",
			feature_layer="pool5",

			classifier_layers=["fc6"],
		)

	@property
	def functions(self):
		return super(BaseResNet, self).functions


class ResNet35(BaseResNet, chainer.Chain):
	n_layers = 35

	def init_extra_layers(self, n_classes, **kwargs):
		self.conv1 = L.Convolution2D(3, 64, 7, 2, 3, **kwargs)
		self.bn1 = L.BatchNormalization(64)
		self.res2 = BuildingBlock(2, 64, 64, 256, 1, **kwargs)
		self.res3 = BuildingBlock(3, 256, 128, 512, 2, **kwargs)
		self.res4 = BuildingBlock(3, 512, 256, 1024, 2, **kwargs)
		self.res5 = BuildingBlock(3, 1024, 512, 2048, 2, **kwargs)
		self.fc6 = L.Linear(2048, n_classes)

	@property
	def functions(self):
		links = [
			("conv1", [self.conv1, self.bn1, F.relu]),
			("pool1", [partial(F.max_pooling_2d, ksize=3, stride=2)]),
			("res2", [self.res2]),
			("res3", [self.res3]),
			("res4", [self.res4]),
			("res5", [self.res5]),
			("pool5", [self.pool]),
			("fc6", [self.fc6]),
			("prob", [F.softmax]),
		]
		return OrderedDict(links)

class ResNet50(BaseResNet, L.ResNet50Layers):
	n_layers = 50


class ResNet101(BaseResNet, L.ResNet101Layers):
	n_layers = 101


class ResNet152(BaseResNet, L.ResNet152Layers):
	n_layers = 152