diff --git a/multifit/training.py b/multifit/training.py index 990d9fa..8309950 100644 --- a/multifit/training.py +++ b/multifit/training.py @@ -167,10 +167,10 @@ class ULMFiTTrainingCommand(Params): f.write(json_str) return json_str - def load_(self, experiment_path, tantetive=True, update_arch=True, silent=False): + def load_(self, experiment_path, tentative=True, update_arch=True, silent=False): fn = experiment_path / self.info_json if not fn.exists(): - if not tantetive: + if not tentative: warn(f"Unable to load experiment_path {experiment_path}") return False print(f"Loading {fn}") @@ -179,7 +179,7 @@ class ULMFiTTrainingCommand(Params): base = d.pop('base', None) arch = d.pop('arch') if hasattr(self, 'base'): - self.base.load_(Path(base), tantetive=True, update_arch=False, silent=silent) + self.base.load_(Path(base), tentative=True, update_arch=False, silent=silent) # compatiblity with older info.json formats where lang was not stored self.name = experiment_path.name # V ./de-1/models/fsp15k/multifit_fp16 -> ./de-1 @@ -263,6 +263,7 @@ class ULMFiTPretraining(ULMFiTTrainingCommand): tokenizer = self.base.tokenizer else: tokenizer = self.arch.new_tokenizer() + assert tokenizer is not None dataset = self._set_dataset_(dataset_or_path, tokenizer) learn = self.get_learner(data_lm=dataset.load_lm_databunch(bs=self.bs, bptt=self.bptt)) @@ -272,10 +273,10 @@ class ULMFiTPretraining(ULMFiTTrainingCommand): self._fit_schedule(learn) self.experiment_path = experiment_path - tokenizer.save(self.experiment_path, learn=learn) learn.to_fp32() learn.save_encoder(ENC_BEST) learn.save(LM_BEST, with_opt=False) + tokenizer.save(self.experiment_path, learn=learn) learn.destroy() self.save_paramters() print("Language model saved to", self.experiment_path) @@ -351,7 +352,7 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): arch: ULMFiTArchitecture = None def get_learner(self, data_clas, eval_only=False, loss_func=None, **additional_trn_args): - assert self.weighted_cross_entropy is None or self.label_smoothing_eps == 0, "Label smoohting not implemented with weighted_cross_entropy" + assert self.weighted_cross_entropy is None or self.label_smoothing_eps == 0, "Label smoothing not implemented with weighted_cross_entropy" if loss_func is None: if self.weighted_cross_entropy is not None: loss_func = CrossEntropyFlat(weight=torch.tensor(self.weighted_cross_entropy, dtype=torch.float32).cuda()) @@ -403,6 +404,7 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): self.replace_(**train_config, _strict=True) base_tokenizer = self.base.tokenizer + assert base_tokenizer dataset = self._set_dataset_(dataset_or_path, base_tokenizer) data_clas = dataset.load_clas_databunch(bs=self.bs, classes=classes, label_cls=label_cls, label_cols=label_cols) learn = self.get_learner(data_clas=data_clas, loss_func=loss_func) @@ -413,15 +415,17 @@ class ULMFiTClassifier(ULMFiTTrainingCommand): print(f"Training: {learn.path / learn.model_dir}") learn.unfreeze() self.experiment_path = learn.path / learn.model_dir - self._fit_schedule(learn) - + try: + self._fit_schedule(learn) + except KeyBoardInterrupt: + pass base_tokenizer.save(self.experiment_path, learn=learn) learn.to_fp32() learn.save(CLS_BEST, with_opt=False) print("Classifier model saved to", self.experiment_path) self.save_paramters() learn.destroy() - return + return learn def _validate(self, learn, ds_type): ds_name = ds_type.name.lower()