Просмотр исходного кода

fixed correct device usage. fixed logging

Dimitri Korsch 6 лет назад
Родитель
Сommit
f0b5045e94
2 измененных файлов с 4 добавлено и 3 удалено
  1. 1 1
      cvfinetune/finetuner/base.py
  2. 3 2
      cvfinetune/parser.py

+ 1 - 1
cvfinetune/finetuner/base.py

@@ -315,6 +315,7 @@ class DefaultFinetuner(_ModelMixin, _DatasetMixin, _TrainerMixin):
 		super(DefaultFinetuner, self).__init__(*args, **kwargs)
 
 		self.gpu_config(opts, *args, **kwargs)
+		cuda.get_device_from_id(self.device).use()
 
 		self.init_annotations(opts)
 		self.init_model(opts)
@@ -334,5 +335,4 @@ class DefaultFinetuner(_ModelMixin, _DatasetMixin, _TrainerMixin):
 			self.device = -1
 		else:
 			self.device = opts.gpu[0]
-		cuda.get_device_from_id(self.device).use()
 

+ 3 - 2
cvfinetune/parser.py

@@ -80,7 +80,7 @@ def default_factory(extra_list=[]):
 
 
 class FineTuneParser(GPUParser):
-	def init_logger(self, simple=False):
+	def init_logger(self, simple=False, logfile=None):
 		if not self.has_logging: return
 		fmt = '{levelname:s} - [{asctime:s}] {filename:s}:{lineno:d} [{funcName:s}]: {message:s}'
 
@@ -88,7 +88,8 @@ class FineTuneParser(GPUParser):
 		handler0.addFilter(HostnameFilter())
 		handler0.setFormatter(logging.Formatter("<{hostname:^10s}>: " + fmt, style="{"))
 
-		handler1 = logging.FileHandler(filename=f"{platform.node()}.log", mode="w")
+		filename = logfile if logfile is not None else f"{platform.node()}.log"
+		handler1 = logging.FileHandler(filename=filename, mode="w")
 		handler1.setFormatter(logging.Formatter(fmt, style="{"))
 
 		logger = logging.getLogger()