diff --git a/demo/image-classification/predict.py b/demo/image-classification/predict.py index 7922663a5f429c30fd7eb2308c730e9e7d069e10..9a1d888592dcab271e63552fb3a1c1eac7ab6fa3 100644 --- a/demo/image-classification/predict.py +++ b/demo/image-classification/predict.py @@ -8,7 +8,7 @@ import numpy as np # yapf: disable parser = argparse.ArgumentParser(__doc__) -parser.add_argument("--use_gpu", type=bool, default=False, help="Whether use GPU for predict.") +parser.add_argument("--use_gpu", type=bool, default=True, help="Whether use GPU for predict.") parser.add_argument("--checkpoint_dir", type=str, default="paddlehub_finetune_ckpt", help="Path to save log data.") parser.add_argument("--batch_size", type=int, default=16, help="Total examples' number in batch for training.") parser.add_argument("--module", type=str, default="resnet50", help="Module used as a feature extractor.") @@ -70,15 +70,17 @@ def predict(args): data = ["./test/test_img_daisy.jpg", "./test/test_img_roses.jpg"] label_map = dataset.label_dict() - for result in task.predict(data=data): - result = np.argmax(result, axis=2) - index = 0 - for batch in result: - for predict_result in batch: - index += 1 - predict_result = label_map[predict_result] - print("input %i is %s, and the predict result is %s" % - (index, data[index - 1], predict_result)) + index = 0 + # get classification result + results = task.predict(data=data) + for batch_result in results: + # get predict index + batch_result = np.argmax(batch_result, axis=2)[0] + for result in batch_result: + index += 1 + result = label_map[result] + print("input %i is %s, and the predict result is %s" % + (index, data[index - 1], result)) if __name__ == "__main__": diff --git a/paddlehub/finetune/checkpoint.py b/paddlehub/finetune/checkpoint.py index 9e436b89eb734599589d4bb98b1b89e4d11b50fe..21bf6a3e16e0b6787e87329806823816ee5cf89a 100644 --- a/paddlehub/finetune/checkpoint.py +++ b/paddlehub/finetune/checkpoint.py @@ -30,31 +30,36 @@ CKPT_FILE_NAME = "ckpt.meta" def load_checkpoint(checkpoint_dir, exe, main_program=fluid.default_main_program(), - startup_program=fluid.default_startup_program()): + startup_program=fluid.default_startup_program(), + load_best_model=False): ckpt_meta_path = os.path.join(checkpoint_dir, CKPT_FILE_NAME) + ckpt = checkpoint_pb2.CheckPoint() logger.info("Try loading checkpoint from {}".format(ckpt_meta_path)) if os.path.exists(ckpt_meta_path): - ckpt = checkpoint_pb2.CheckPoint() with open(ckpt_meta_path, "rb") as f: ckpt.ParseFromString(f.read()) + current_epoch = 1 + global_step = 0 + best_model_path = os.path.join(checkpoint_dir, "best_model") + if load_best_model and os.path.exists(best_model_path): + fluid.io.load_persistables(exe, best_model_path, main_program) + logger.info("PaddleHub model best model loaded.") + return current_epoch, global_step + elif ckpt.latest_model_dir: fluid.io.load_persistables(exe, ckpt.latest_model_dir, main_program) logger.info("PaddleHub model checkpoint loaded. current_epoch={}, " "global_step={}".format(ckpt.current_epoch, ckpt.global_step)) return ckpt.current_epoch, ckpt.global_step - else: - current_epoch = 1 - global_step = 0 - latest_model_dir = None - logger.info( - "PaddleHub model checkpoint not found, start training from scratch..." - ) - exe.run(startup_program) - return current_epoch, global_step + logger.info( + "PaddleHub model checkpoint not found, start training from scratch...") + exe.run(startup_program) + + return current_epoch, global_step def save_checkpoint(checkpoint_dir, diff --git a/paddlehub/finetune/task.py b/paddlehub/finetune/task.py index 1c4d2445a17339b32c1324448433f7e832bb75ed..879ca1ef9db6ca9c0d964222c3baec2509f41fdb 100644 --- a/paddlehub/finetune/task.py +++ b/paddlehub/finetune/task.py @@ -128,10 +128,11 @@ class BasicTask(object): # run environment self._phases = [] self._envs = {} + self._predict_data = None - def init_if_necessary(self): + def init_if_necessary(self, load_best_model=False): if not self._load_checkpoint: - self.load_checkpoint() + self.load_checkpoint(load_best_model=load_best_model) self._load_checkpoint = True @contextlib.contextmanager @@ -159,6 +160,11 @@ class BasicTask(object): self.env.loss = self._add_loss() self.env.metrics = self._add_metrics() + if self.is_predict_phase or self.is_test_phase: + self.env.main_program = self.env.main_program.clone(for_test=True) + hub.common.paddle_helper.set_op_attr( + self.env.main_program, is_test=True) + if self.config.use_pyreader: t_program = fluid.Program() with fluid.program_guard(t_program, self.env.startup_program): @@ -291,8 +297,12 @@ class BasicTask(object): @property def reader(self): + if self.is_predict_phase: + data = self._predict_data + else: + data = None self.env.reader = self._base_data_reader.data_generator( - batch_size=self.config.batch_size, phase=self.phase) + batch_size=self.config.batch_size, phase=self.phase, data=data) return self.env.reader @property @@ -315,8 +325,6 @@ class BasicTask(object): @property def output(self): - if self.is_predict_phase: - raise RuntimeError() if not self.env.is_inititalized: self._build_env() return self.env.output @@ -412,7 +420,8 @@ class BasicTask(object): self.config.checkpoint_dir, self.exe, main_program=self.main_program, - startup_program=self._base_startup_program) + startup_program=self._base_startup_program, + load_best_model=load_best_model) if load_best_model: model_saved_dir = os.path.join(self.config.checkpoint_dir, @@ -454,10 +463,12 @@ class BasicTask(object): self._eval_end_event(run_states) def predict(self, data, load_best_model=True): - with self.phase_guard(phase=phase): - self.init_if_necessary() - for run_state in self._run(): - yield run_state.run_results + with self.phase_guard(phase="predict"): + self._predict_data = data + self.init_if_necessary(load_best_model=load_best_model) + run_states = self._run() + self._predict_data = None + return [run_state.run_results for run_state in run_states] def _run(self, do_eval=False): with fluid.program_guard(self.main_program, self.startup_program):