MCPcopy Create free account
hub / github.com/CHENGY12/PLOT / load_model

Method load_model

plot-pp/trainers/plotpp.py:365–396  ·  view source on GitHub ↗
(self, directory, epoch=None)

Source from the content-addressed store, hash-verified

363
364
365 def load_model(self, directory, epoch=None):
366 if not directory:
367 print("Note that load_model() is skipped as no pretrained model is given")
368 return
369
370 names = self.get_model_names()
371
372 # By default, the best model is loaded
373 model_file = "model-best.pth.tar"
374
375 if epoch is not None:
376 model_file = "model.pth.tar-" + str(epoch)
377
378 for name in names:
379 model_path = osp.join(directory, name, model_file)
380
381 if not osp.exists(model_path):
382 raise FileNotFoundError('Model not found at "{}"'.format(model_path))
383
384 checkpoint = load_checkpoint(model_path)
385 state_dict = checkpoint["state_dict"]
386 epoch = checkpoint["epoch"]
387
388 # Ignore fixed token vectors
389 if "prompt_learner.token_prefix" in state_dict:
390 del state_dict["prompt_learner.token_prefix"]
391
392 if "prompt_learner.token_suffix" in state_dict:
393 del state_dict["prompt_learner.token_suffix"]
394
395 print("Loading weights to {} " 'from "{}" (epoch = {})'.format(name, model_path, epoch))
396 self._models[name].load_state_dict(state_dict, strict=False)
397
398
399

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected