瀏覽代碼

updated handling of kwargs in the classifier creation logic

Dimitri Korsch 3 年之前
父節點
當前提交
f030b6a761
共有 1 個文件被更改,包括 6 次插入5 次删除
  1. 6 5
      cvfinetune/finetuner/mixins/classifier.py

+ 6 - 5
cvfinetune/finetuner/mixins/classifier.py

@@ -16,8 +16,8 @@ class _ClassifierCreator:
         self.kwargs = kwargs
 
     def __call__(self, *args, **kwargs):
-        kwargs = dict(self.kwargs, **kwargs)
-        return self.cls(*args, **kwargs)
+        self.kwargs = dict(self.kwargs, **kwargs)
+        return self.cls(*args, **self.kwargs)
 
 class _ClassifierMixin(BaseMixin):
     """
@@ -39,12 +39,13 @@ class _ClassifierMixin(BaseMixin):
         self._label_smoothing = label_smoothing
 
 
-    def init_classifier(self):
+    def init_classifier(self, **kwargs):
         self._check_attr("model")
         self._check_attr("n_classes")
 
-        self.clf = self._clf_creator(model=self.model,
-                                     loss_func=self.loss_func)
+        self.clf = self._clf_creator(self.model,
+                                     loss_func=self.loss_func,
+                                     **kwargs)
 
         kwargs = self._clf_creator.kwargs
         logging.info(