some changes in model loading

cvfinetune/__init__.py

@@ -1 +1 @@
-__version__ = "0.5.5"
+__version__ = "0.6.0"

cvfinetune/classifier.py

@@ -1,78 +0,0 @@
-import abc
-import chainer
-import chainer.functions as F
-import chainer.links as L
-import logging
-
-from chainer_addons.models.classifier import Classifier as C
-
-class Classifier(C):
-
-	def __init__(self, *args, **kwargs):
-		super(Classifier, self).__init__(*args, **kwargs)
-
-		assert hasattr(self, "model"), \
-			"This classifiert has no \"model\" attribute!"
-
-	@property
-	def feat_size(self):
-		if hasattr(self.model.pool, "output_dim") and self.model.pool.output_dim is not None:
-			return self.model.pool.output_dim
-
-		return self.model.meta.feature_size
-
-	@property
-	def output_size(self):
-		return self.feat_size
-
-	def enable_only_head(self):
-		self.model.disable_update()
-		self.model.fc.enable_update()
-
-
-class SeparateModelClassifier(Classifier):
-	"""Classifier, that holds two separate models"""
-
-	def __init__(self, *args, **kwargs):
-		super(SeparateModelClassifier, self).__init__(*args, **kwargs)
-
-		with self.init_scope():
-			self.init_separate_model()
-
-	@abc.abstractmethod
-	def __call__(self, *args, **kwargs):
-		super(SeparateModelClassifier, self).__call__(*args, **kwargs)
-
-	def init_separate_model(self):
-
-		if hasattr(self, "separate_model"):
-			logging.warn("Global Model already initialized! Skipping further execution!")
-			return
-
-		self.separate_model = self.model.copy(mode="copy")
-
-	def loader(self, model_loader):
-
-		def inner(n_classes, feat_size):
-			# use the given feature size here
-			model_loader(n_classes=n_classes, feat_size=feat_size)
-
-			# use the given feature size first ...
-			self.separate_model.reinitialize_clf(
-				n_classes=n_classes,
-				feat_size=feat_size)
-
-			# then copy model params ...
-			self.separate_model.copyparams(self.model)
-
-			# now use the default feature size to re-init the classifier
-			self.separate_model.reinitialize_clf(
-				n_classes=n_classes,
-				feat_size=self.feat_size)
-
-		return inner
-
-	def enable_only_head(self):
-		super(SeparateModelClassifier, self).enable_only_head()
-		self.separate_model.disable_update()
-		self.separate_model.fc.enable_update()

cvfinetune/finetuner/base.py

@@ -35,7 +35,7 @@ class DefaultFinetuner(mixins._ModelMixin, mixins._DatasetMixin, mixins._Trainer
 		self.init_iterators(opts)
 
 		self.init_classifier(opts)
-		self.load_model_weights(opts)
+		self.load_weights(opts)
 
 		self.init_optimizer(opts)
 		self.init_updater()

cvfinetune/finetuner/mixins/model.py

@@ -10,8 +10,11 @@ from chainer_addons.training import optimizer
 from chainer_addons.training import optimizer_hooks
 from cvdatasets.dataset.image import Size
 from cvdatasets.utils import pretty_print_dict
+from cvmodelz.models import ModelFactory
 from functools import partial
 from pathlib import Path
+from typing import Callable
+from typing import Tuple
 
 
 class _ModelMixin(abc.ABC):
@@ -32,6 +35,14 @@ class _ModelMixin(abc.ABC):
 	def model_info(self):
 		return self.data_info.MODELS[self.model_type]
 
+	def init_model(self, opts):
+		"""creates backbone CNN model. This model is wrapped around the classifier later"""
+
+		self.model = ModelFactory.new(self.model_type,
+			input_size=Size(opts.input_size),
+			**self.model_kwargs
+		)
+
 	def init_classifier(self, opts):
 
 		clf_class, kwargs = self.classifier_cls, self.classifier_kwargs
@@ -97,63 +108,59 @@ class _ModelMixin(abc.ABC):
 			logging.warning("========= Fine-tuning only classifier layer! =========")
 			enable_only_head(self.clf)
 
-	def init_model(self, opts):
-		"""creates backbone CNN model. This model is wrapped around the classifier later"""
 
-		if self.model_type.startswith("cv2_"):
-			model_type = args.model_type.split("cv2_")[-1]
+	def _get_loader(self, opts) -> Tuple[bool, str]:
+		if getattr(opts, "from_scratch", False):
+			logging.info("Training a {0.__class__.__name__} model from scratch!".format(self.model))
+			return None, None
+
+		if getattr(opts, "load", None):
+			weights = getattr(opts, "load", None)
+			logging.info(f"Loading already fine-tuned weights from \"{weights}\"")
+			return False, weights
+
+		elif getattr(opts, "weights", None):
+			weights = getattr(opts, "weights", None)
+			logging.info(f"Loading custom fine-tuned weights from \"{weights}\"")
+			return True, weights
+
 		else:
-			model_type = self.model_info.class_key
+			weights = self._default_weights(opts)
+			logging.info(f"Loading custom fine-tuned weights from \"{weights}\"")
+			return True, weights
 
-		self.model = ModelType.new(
-			model_type=model_type,
-			input_size=Size(opts.input_size),
-			**self.model_kwargs,
+	def _default_weights(self, opts):
+		ds_info = self.data_info
+		model_info = self.model_info
+
+		base_dir = Path(ds_info.BASE_DIR)
+		weights_dir = base_dir / ds_info.MODEL_DIR / model_info.folder
+
+		weights = model_info.weights
+		### TODO: make pre-training command line argument!
+		return str(weights_dir / weights.get("inat", weights.get("imagenet")))
+
+
+	def load_weights(self, opts) -> None:
+
+		finetune, weights = self._get_loader(opts)
+
+		self.clf.load(weights,
+			n_classes=self.n_classes,
+			finetune=finetune,
+
+			path=opts.load_path,
+			strict=opts.load_strict,
+			headless=opts.headless
 		)
 
-	def load_model_weights(self, args):
-		if getattr(args, "from_scratch", False):
-			logging.info("Training a {0.__class__.__name__} model from scratch!".format(self.model))
-			loader = self.model.reinitialize_clf
-			self.weights = None
-		else:
-			if args.load:
-				self.weights = args.load
-				msg = "Loading already fine-tuned weights from \"{}\""
-				loader_func = self.model.load_for_inference
-			else:
-				if args.weights:
-					msg = "Loading custom pre-trained weights \"{}\""
-					self.weights = args.weights
-
-				else:
-					msg = "Loading default pre-trained weights \"{}\""
-					self.weights = str(Path(
-						self.data_info.BASE_DIR,
-						self.data_info.MODEL_DIR,
-						self.model_info.folder,
-						self.model_info.weights
-					))
-
-				loader_func = self.model.load_for_finetune
-
-			logging.info(msg.format(self.weights))
-			kwargs = dict(
-				weights=self.weights,
-				strict=args.load_strict,
-				path=args.load_path,
-				headless=args.headless,
-			)
-			loader = partial(loader_func, **kwargs)
+		self.clf.cleargrads()
 
 		feat_size = self.model.meta.feature_size
 
 		if hasattr(self.clf, "output_size"):
 			feat_size = self.clf.output_size
 
-		if hasattr(self.clf, "loader"):
-			loader = self.clf.loader(loader)
+		### TODO: handle feature size!
 
 		logging.info(f"Part features size after encoding: {feat_size}")
-		loader(n_classes=self.n_classes, feat_size=feat_size)
-		self.clf.cleargrads()

cvfinetune/parser/dataset_args.py

@@ -1,12 +1,13 @@
 import abc
 
 from cvargparse import Arg
+from cvargparse import BaseParser
 from cvfinetune.parser.utils import DEFAULT_INFO_FILE
 from cvfinetune.parser.utils import get_info_file
 from cvfinetune.parser.utils import parser_extender
 
 @parser_extender
-def add_dataset_args(parser):
+def add_dataset_args(parser: BaseParser) -> None:
 
 	info_file = get_info_file()
 

cvfinetune/parser/model_args.py

@@ -1,33 +1,20 @@
 import abc
 
-from cvargparse import Arg
-from cvfinetune.parser.utils import get_info_file
-from cvfinetune.parser.utils import parser_extender
-
 from chainer_addons.links import PoolingType
 from chainer_addons.models import PrepareType
-
-class ModelChoices(object):
-
-	def __init__(self, choices=[]):
-		super(ModelChoices, self).__init__()
-		self.choices = choices
-
-	def __contains__(self, value):
-		return value.startswith("cv2_") or value in self.choices
-
-	def __iter__(self):
-		return iter(self.choices + ["cv2_<any other model>"])
+from cvargparse import Arg
+from cvargparse import BaseParser
+from cvfinetune.parser.utils import parser_extender
+from cvmodelz.models import ModelFactory
 
 @parser_extender
-def add_model_args(parser):
-
-	info_file = get_info_file()
-	choices = None if info_file is None else info_file.MODELS.keys()
+def add_model_args(parser: BaseParser) -> None:
 
+	choices = ModelFactory.get_models(["chainercv2", "cvmodelz"])
 	_args = [
 		Arg("--model_type", "-mt",
-			default="cv2_resnet50", choices=ModelChoices(choices),
+			required=True,
+			choices=choices,
 			help="type of the model"),
 
 		Arg("--input_size", type=int, nargs="+", default=0,
@@ -54,7 +41,7 @@ def add_model_args(parser):
 		Arg("--load_strict", action="store_true",
 			help="load weights in a strict mode"),
 
-		Arg("--load_path", type=str,
+		Arg("--load_path", type=str, default="",
 			help="load path within the weights archive"),
 	]
 

cvfinetune/parser/training_args.py

@@ -2,12 +2,13 @@ import abc
 
 from cvargparse import Arg
 from cvargparse import ArgFactory
+from cvargparse import BaseParser
 from cvfinetune.parser.utils import parser_extender
 
 from chainer_addons.training import OptimizerType
 
 @parser_extender
-def add_training_args(parser):
+def add_training_args(parser: BaseParser) -> None:
 
 	_args = ArgFactory([