Bläddra i källkod

copied and adapted content from chainer_addons

Dimitri Korsch 4 år sedan
förälder
incheckning
f58798c481

+ 11 - 0
.editorconfig

@@ -0,0 +1,11 @@
+# top-most EditorConfig file
+root = true
+
+[*]
+end_of_line = lf
+insert_final_newline = true
+
+[*.py]
+indent_style = tab
+indent_size = 4
+charset = utf-8

+ 1 - 1
Makefile

@@ -15,4 +15,4 @@ get_version:
 	@python setup.py --version
 
 run_tests:
-	bash scripts/tests.sh
+	python tests/main.py

+ 1 - 1
README.md

@@ -1,3 +1,3 @@
 # cvmodelz
 
-*Why is it written with 'z'? Because 'cvmodels' already exists ¯\_(ツ)_/¯*
+*"Why is it written with 'z'? Because 'cvmodels' already exists ¯\\\_(ツ)\_/¯"*

+ 10 - 0
cvmodelz/__init__.py

@@ -0,0 +1,10 @@
+from cvmodelz.models.meta_info import ModelInfo
+from cvmodelz.models.wrapper import ModelWrapper
+from cvmodelz.models.pretrained import PretrainedModelMixin
+
+
+__all__ = [
+	"ModelInfo",
+	"ModelWrapper",
+	"PretrainedModelMixin",
+]

+ 59 - 0
cvmodelz/models/__init__.py

@@ -0,0 +1,59 @@
+from chainer import links as L
+
+from cvmodelz.models.base import BaseModel
+from cvmodelz.models import pretrained
+
+__all__ = [
+	"BaseModel"
+]
+
+
+supported = dict(
+	chainer=(
+		L.ResNet50Layers,
+		L.ResNet101Layers,
+		L.ResNet152Layers,
+		L.VGG16Layers,
+		L.VGG19Layers,
+	),
+
+	chainercv=(),
+
+	chainercv2=(),
+
+	custom=(
+		pretrained.VGG16,
+		pretrained.VGG19,
+
+		pretrained.ResNet35,
+		pretrained.ResNet50,
+		pretrained.ResNet101,
+		pretrained.ResNet152,
+	),
+)
+
+def _check(model, key):
+	global supported
+	return isinstance(model, supported[key])
+
+
+def is_chainer_model(model):
+	return _check(model, "chainer")
+
+def is_cv_model(model):
+	return _check(model, "chainercv")
+
+def is_cv2_model(model):
+	return _check(model, "chainercv2")
+
+def is_custom_model(model):
+	return _check(model, "custom")
+
+
+if __name__ == '__main__':
+	from cvmodelz import utils
+
+	# model = L.VGG19Layers(pretrained_model=None)
+	model = pretrained.ResNet35()
+	utils.print_model_info(model, input_size=224)
+

+ 73 - 0
cvmodelz/models/base.py

@@ -0,0 +1,73 @@
+import abc
+import chainer.functions as F
+import chainer.links as L
+
+from cvmodelz import utils
+
+class BaseModel(abc.ABC):
+
+	@abc.abstractmethod
+	def __call__(self, X, layer_name=None):
+		pass
+
+	@abc.abstractproperty
+	def model_instance(self):
+		raise NotImplementedError()
+
+	@property
+	def clf_layer_name(self):
+		return self.meta.classifier_layers[-1]
+
+	@property
+	def clf_layer(self):
+		return utils.get_attr_from_path(self.model_instance, self.clf_layer_name)
+
+	def loss(self, pred, gt, loss_func=F.softmax_cross_entropy):
+		return loss_func(pred, gt)
+
+	def accuracy(self, pred, gt):
+		return F.accuracy(pred, gt)
+
+	def reinitialize_clf(self, n_classes, feat_size=None, initializer=None):
+
+		if initializer is None or not callable(initializer):
+			initializer = HeNormal(scale=1.0)
+
+		clf_layer = self.clf_layer
+
+		assert isinstance(clf_layer, L.Linear)
+
+		w_shape = (n_classes, feat_size or clf_layer.W.shape[1])
+		dtype = clf_layer.W.dtype
+
+		clf_layer.W.data = np.zeros(w_shape, dtype=dtype)
+		clf_layer.b.data = np.zeros(w_shape[0], dtype=dtype)
+
+		initializer(clf_layer.W.data)
+
+	def load_for_finetune(self, weights, n_classes, *, path="", strict=False, headless=False, **kwargs):
+		"""
+			The weights should be pre-trained on a bigger
+			dataset (eg. ImageNet). The classification layer is
+			reinitialized after all other weights are loaded
+		"""
+		self.load(weights, path=path, strict=strict, headless=headless)
+		self.reinitialize_clf(n_classes, **kwargs)
+
+	def load_for_inference(self, weights, n_classes, *, path="", strict=False, headless=False, **kwargs):
+		"""
+			In this use case we are loading already fine-tuned
+			weights. This means, we need to reinitialize the
+			classification layer first and then load the weights.
+		"""
+		self.reinitialize_clf(n_classes, **kwargs)
+		self.load(weights, path=path, strict=strict, headless=headless)
+
+	def load(self, weights, *, path="", strict=False, headless=False):
+		if weights not in [None, "auto"]:
+			ignore_names = None
+			if headless:
+				ignore_names = lambda name: name.startswith(path + self.clf_layer_name)
+
+			npz.load_npz(weights, self.model_instance,
+				path=path, strict=strict, ignore_names=ignore_names)

+ 26 - 0
cvmodelz/models/meta_info.py

@@ -0,0 +1,26 @@
+import pyaml
+
+from dataclasses import dataclass
+from typing import Tuple
+
+
+@dataclass
+class ModelInfo(object):
+
+	name:                       str         = "GenericModel"
+	input_size:                 int         = 224
+	feature_size:               int         = 2048
+	n_conv_maps:                int         = 2048
+
+	conv_map_layer:             str         = "conv"
+	feature_layer:              str         = "fc"
+
+	classifier_layers:          Tuple[str]  = ("fc",)
+
+	def __str__(self):
+		obj = dict(ModelInfo=self.__dict__)
+		return pyaml.dump(obj, sort_dicts=False, )
+
+
+if __name__ == '__main__':
+	print(ModelInfo())

+ 18 - 0
cvmodelz/models/pretrained/__init__.py

@@ -0,0 +1,18 @@
+from cvmodelz.models.pretrained.base import PretrainedModelMixin
+from cvmodelz.models.pretrained.vgg import VGG16
+from cvmodelz.models.pretrained.vgg import VGG19
+from cvmodelz.models.pretrained.resnet import ResNet35
+from cvmodelz.models.pretrained.resnet import ResNet50
+from cvmodelz.models.pretrained.resnet import ResNet101
+from cvmodelz.models.pretrained.resnet import ResNet152
+
+
+__all__ = [
+	"PretrainedModelMixin",
+	"VGG16",
+	"VGG19",
+	"ResNet35",
+	"ResNet50",
+	"ResNet101",
+	"ResNet152",
+]

+ 61 - 0
cvmodelz/models/pretrained/base.py

@@ -0,0 +1,61 @@
+import abc
+
+from chainer import functions as F
+from typing import Callable
+
+from cvmodelz import models
+from cvmodelz.models.base import BaseModel
+from cvmodelz.models.meta_info import ModelInfo
+
+class PretrainedModelMixin(BaseModel):
+	"""
+		This mixin is designed to be a superclass besides one of
+		chainer's built in models (VGG, ResNet, GoogLeNet).
+
+		Example:
+
+			import chainer.links as L
+			class OurResNet(PretrainedModelMixin, L.ResNet50layers):
+				...
+	"""
+
+	def __init__(self, n_classes: int = 1000, pooling: Callable = F.identity, *args, **kwargs):
+
+		if models.is_chainer_model(self):
+			kwargs["pretrained_model"] = None
+
+		super(PretrainedModelMixin, self).__init__(*args, **kwargs)
+
+		with self.init_scope():
+			self.init_extra_layers(n_classes)
+			self.pool = pooling
+
+	def __call__(self, X, layer_name=None):
+		assert hasattr(self, "meta"), "Did you forgot to initialize the meta attribute?"
+
+		layer_name = layer_name or self.meta.classifier_layers[-1]
+		caller = super(PretrainedModelMixin, self).__call__
+		activations = caller(X, layers=[layer_name])
+
+		if isinstance(activations, dict):
+			activations = activations[layer_name]
+
+		return activations
+
+	def init_extra_layers(self, *args, **kwargs):
+		pass
+
+	# @abc.abstractproperty
+	# def _links(self):
+	# 	raise NotImplementedError()
+
+	# @property
+	# def functions(self):
+	# 	return OrderedDict(self._links)
+
+	@property
+	def model_instance(self):
+		""" since it is a mixin, we are the model """
+
+		return self
+

+ 66 - 0
cvmodelz/models/pretrained/resnet.py

@@ -0,0 +1,66 @@
+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

+ 33 - 0
cvmodelz/models/pretrained/vgg.py

@@ -0,0 +1,33 @@
+from chainer import links as L
+from chainer.links.model.vision.vgg import prepare
+from chainer.links.model.vision.vgg import _max_pooling_2d
+
+from cvmodelz.models.meta_info import ModelInfo
+from cvmodelz.models.pretrained.base import PretrainedModelMixin
+
+def _vgg_meta(final_conv_layer):
+	return ModelInfo(
+		name="VGG",
+		input_size=224,
+		feature_size=4096,
+		n_conv_maps=512,
+
+		conv_map_layer=final_conv_layer,
+		feature_layer="fc7",
+
+		classifier_layers=["fc6", "fc7", "fc8"],
+	)
+
+
+class VGG19(PretrainedModelMixin, L.VGG19Layers):
+
+	def __init__(self, *args, **kwargs):
+		super(VGG19, self).__init__(*args, pooling=_max_pooling_2d, **kwargs)
+		self.meta = _vgg_meta("conv5_3")
+
+class VGG16(PretrainedModelMixin, L.VGG16Layers):
+
+	def __init__(self, *args, **kwargs):
+		super(VGG16, self).__init__(*args, pooling=_max_pooling_2d, **kwargs)
+		self.meta = _vgg_meta("conv5_4")
+

+ 67 - 0
cvmodelz/models/wrapper.py

@@ -0,0 +1,67 @@
+import chainer
+
+from chainer import functions as F
+from typing import Callable
+
+from cvmodelz.models.base import BaseModel
+from cvmodelz.models.meta_info import ModelInfo
+
+class ModelWrapper(BaseModel, chainer.Chain):
+	"""
+		This class is designed to wrap around chainercv2 models
+		and provide the loading API of the BaseModel class.
+		The wrapped model is stored under self.wrapped
+	"""
+
+	def __init__(self, model: chainer.Chain, pooling: Callable = F.identity):
+		super(ModelWrapper, self).__init__()
+
+		name = model.__class__.__name__
+		self.__class__.__name__ = name
+
+		if hasattr(model, "meta"):
+			self.meta = model.meta
+
+		else:
+			self.meta = ModelInfo(
+				name=name,
+				classifier_layers=("output/fc",),
+				conv_map_layer="stage4",
+				feature_layer="pool",
+			)
+
+		with self.init_scope():
+			self.wrapped = model
+			self.pool = pooling
+			delattr(self.wrapped.features, "final_pool")
+
+		self.meta.feature_size = self.clf_layer.W.shape[-1]
+
+	@property
+	def model_instance(self):
+		return self.wrapped
+
+	def load_for_inference(self, *args, path="", **kwargs):
+		return super(ModelWrapper, self).load_for_inference(*args, path=f"{path}wrapped/", **kwargs)
+
+	def __call__(self, X, layer_name=None):
+		if layer_name is None:
+			res = self.wrapped(X)
+
+		elif layer_name == self.meta.conv_map_layer:
+			res = self.wrapped.features(X)
+
+		elif layer_name == self.meta.feature_layer:
+			conv = self.wrapped.features(X)
+			res = self.pool(conv)
+
+		elif layer_name == self.clf_layer_name:
+			conv = self.wrapped.features(X)
+			feat = self.pool(conv)
+			res = self.wrapped.output(feat)
+
+		else:
+			raise ValueError(f"Dont know how to compute \"{layer_name}\"!")
+
+		return res
+

+ 91 - 0
cvmodelz/utils/__init__.py

@@ -0,0 +1,91 @@
+from functools import reduce
+import chainer
+import numpy as np
+import sys
+import logging
+
+from functools import partial
+from tabulate import tabulate
+
+def get_attr_from_path(obj, path, *, sep="/"):
+	def getter(o, attr):
+		return
+	getter = lambda o, attr: (getattr(o, attr) if attr else o)
+	return reduce(getter, path.split(sep), obj)
+
+def _get_activation_shapes(model, input_size, input_var, batch_size=2, n_channels=3):
+	assert hasattr(model, "functions"), "Model should have functions defined!"
+	if input_var is None:
+		input_shape = (batch_size, n_channels, input_size, input_size)
+		x = model.xp.zeros(input_shape, dtype=model.xp.float32)
+	else:
+		x = input_var
+
+	res = [("Input", x.shape)]
+	with chainer.no_backprop_mode(), chainer.using_config("train", False):
+		for name, link in model.functions.items():
+			in_shape = str(x.shape)
+			out_shapes = []
+			for func in link:
+				x = func(x)
+				out_shapes.append(str(x.shape))
+			logging.debug("\t".join(map(str, (name, in_shape, out_shapes))))
+			res.append((name, in_shape, " -> ".join(out_shapes)))
+	return res
+
+def print_model_info(model, file=sys.stdout, input_size=None, input_var=None):
+	_print = partial(print, file=file)
+	name = getattr(model, "name", None)
+	if name is None:
+		name = model.__class__.__name__
+
+
+	rows = []
+	default_size = 224
+	if hasattr(model, "meta"):
+		_print(model.meta)
+		default_size = model.meta.input_size
+		rows.append(("Default input size", f"{default_size}"))
+
+		feature_size = model.meta.feature_size
+		rows.append(("Feature size", f"{feature_size}"))
+
+		n_conv_maps = model.meta.n_conv_maps
+		rows.append(("# of conv maps (last layer)", f"{n_conv_maps}"))
+
+	n_weights = model.count_params()
+	rows.append(("# of parameters", f"{n_weights:,d}"))
+
+	n_params = len(list(model.params()))
+	rows.append(("# of trainables", f"{n_params:,d}"))
+
+	n_layers = len(list(model.links()))
+	rows.append(("# of layers", f"{n_layers:,d}"))
+
+	_print(f"Printing some information about \"{name}\" model")
+	_print(tabulate(rows, tablefmt="fancy_grid"))
+
+	shapes = _get_activation_shapes(model, input_size or default_size, input_var)
+	_print("In/Out activation shapes:")
+	_print(tabulate(shapes,
+		headers=["Link name", "Input", "Output"],
+		tablefmt="fancy_grid"))
+
+
+if __name__ == "__main__":
+	class Foo:
+		attr = None
+		def __str__(self):
+			return f"<{hex(id(self))}: attr={self.attr}>"
+
+
+	o = Foo()
+	o.attr = Foo()
+	o.attr.attr = Foo()
+
+	print(o,
+		get_attr_from_path(o, "attr"),
+		get_attr_from_path(o, "attr/"),
+		get_attr_from_path(o, "attr/attr"),
+		sep="\n"
+	)

+ 2 - 0
requirements.txt

@@ -2,3 +2,5 @@
 chainer~=7.0
 chainercv~=0.13
 chainercv2~=0.0
+
+pyaml~=20.4

+ 2 - 1
setup.py

@@ -21,12 +21,13 @@ install_requires = [
 setup(
 	name=pkg_name,
 	version=__version__,
+	python_requires=">3.6",
 	description='Wrapper for various computer vision models (mostly provided by chainer, chainercv, and chainercv2)',
 	log_description=open(str(cwd / "README.md")).read(),
 	author='Dimitri Korsch',
 	author_email='korschdima@gmail.com',
 	license='MIT License',
-	packages=find_packages(),
+	packages=find_packages(exclude=("tests",)),
 	zip_safe=False,
 	setup_requires=[],
 	install_requires=install_requires,

+ 14 - 0
tests/main.py

@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+if __name__ != '__main__': raise Exception("Do not import me!")
+
+import unittest
+import chainer
+import sys
+
+from pathlib import Path
+
+cwd = Path(__file__).resolve().parent
+sys.path.insert(0, str(cwd.parent))
+
+with chainer.using_config("train", False):
+	unittest.main()