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"],
		)


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