def run_benchmark(model): """Measures training time for given model with random data. Args: model: string, file name of the input model. Returns: float, training time in ms. """ engine = ImprintingEngine(test_utils.test_data_path(model), keep_classes=False) extractor = make_interpreter(engine.serialize_extractor_model(), device=':0') extractor.allocate_tensors() width, height = common.input_size(extractor) np.random.seed(12345) # 10 Categories, each has 20 images. data_by_category = collections.defaultdict(list) for i in range(10): for _ in range(20): data_by_category[i].append( np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)) delegate = load_edgetpu_delegate({'device': ':0'}) inference_time = 0. for class_id, tensors in enumerate(data_by_category.values()): for tensor in tensors: common.set_input(extractor, tensor) extractor.invoke() engine.train(classify.get_scores(extractor), class_id=class_id) start = time.perf_counter() interpreter = tflite.Interpreter( model_content=engine.serialize_model(), experimental_delegates=[delegate]) interpreter.allocate_tensors() common.set_input(interpreter, tensors[0]) interpreter.invoke() classify.get_classes(interpreter, top_k=3) inference_time += (time.perf_counter() - start) * 1000 print('Model: %s' % model) print('Inference time: %.2fms' % inference_time) return inference_time
def run_benchmark(model): """Measures training time for given model with random data. Args: model: string, file name of the input model. Returns: float, training time in ms. """ engine = ImprintingEngine( test_utils.test_data_path(model), keep_classes=False) extractor = make_interpreter(engine.serialize_extractor_model()) extractor.allocate_tensors() width, height = common.input_size(extractor) np.random.seed(12345) # 10 Categories, each has 20 images. data_by_category = collections.defaultdict(list) for i in range(10): for _ in range(20): data_by_category[i].append( np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)) start = time.perf_counter() for class_id, tensors in enumerate(data_by_category.values()): for tensor in tensors: common.set_input(extractor, tensor) extractor.invoke() engine.train(classify.get_scores(extractor), class_id=class_id) engine.serialize_model() training_time = (time.perf_counter() - start) * 1000 print('Model: %s' % model) print('Training time: %.2fms' % training_time) return training_time
def _train_and_test(self, model_path, train_points, test_points, keep_classes): # Train. engine = ImprintingEngine(model_path, keep_classes) extractor = make_interpreter(engine.serialize_extractor_model(), device=':0') extractor.allocate_tensors() for point in train_points: for image in point.images: with test_utils.test_image('imprinting', image) as img: set_input(extractor, img) extractor.invoke() embedding = classify.get_scores(extractor) self.assertEqual(len(embedding), engine.embedding_dim) engine.train(embedding, point.class_id) # Test. trained_model = engine.serialize_model() classifier = make_interpreter(trained_model, device=':0') classifier.allocate_tensors() self.assertEqual(len(classifier.get_output_details()), 1) if not keep_classes: self.assertEqual(len(train_points), classify.num_classes(classifier)) for point in test_points: with test_utils.test_image('imprinting', point.image) as img: set_input(classifier, img) classifier.invoke() top = classify.get_classes(classifier, top_k=1)[0] self.assertEqual(top.id, point.class_id) self.assertGreater(top.score, point.score) return trained_model
def train(capture_dir, labels, model, out_model): engine = ImprintingEngine(model, keep_classes=False) extractor = make_interpreter(engine.serialize_extractor_model(), device=':0') extractor.allocate_tensors() for class_id in sorted(labels): class_name = labels[class_id] print('\nClass: %s (id=%d)' % (class_name, class_id)) class_capture_dir = os.path.join(capture_dir, class_name) for img in os.listdir(class_capture_dir): imgpath = os.path.join(class_capture_dir, img) common.set_input(extractor, read_image(imgpath, common.input_size(extractor))) extractor.invoke() embedding = classify.get_scores(extractor) print(' %s => %s' % (imgpath, embedding)) engine.train(embedding, class_id) with open(out_model, 'wb') as f: f.write(engine.serialize_model()) print('\nTrained model was saved to %s' % out_model)
def _transfer_learn_and_evaluate(self, model_path, keep_classes, dataset_path, test_ratio, top_k_range): """Transfer-learns with given params and returns the evaluation result. Args: model_path: string, path of the base model. keep_classes: bool, whether to keep base model classes. dataset_path: string, path to the directory of dataset. The images should be put under sub-directory named by category. test_ratio: float, the ratio of images used for test. top_k_range: int, top_k range to be evaluated. The function will return accuracy from top 1 to top k. Returns: list of float numbers. """ engine = ImprintingEngine(model_path, keep_classes) extractor = make_interpreter(engine.serialize_extractor_model()) extractor.allocate_tensors() num_classes = engine.num_classes print('--------------- Parsing dataset ----------------') print('Dataset path:', dataset_path) # train in fixed order to ensure the same evaluation result. train_set, test_set = test_utils.prepare_data_set_from_directory( dataset_path, test_ratio, True) print('Image list successfully parsed! Number of Categories = ', len(train_set)) print('--------------- Processing training data ----------------') print('This process may take more than 30 seconds.') train_input = [] labels_map = {} for class_id, (category, image_list) in enumerate(train_set.items()): print('Processing {} ({} images)'.format(category, len(image_list))) train_input.append( [os.path.join(dataset_path, category, image) for image in image_list]) labels_map[num_classes + class_id] = category # train print('---------------- Start training -----------------') size = common.input_size(extractor) for class_id, images in enumerate(train_input): for image in images: with test_image(image) as img: common.set_input(extractor, img.resize(size, Image.NEAREST)) extractor.invoke() engine.train(classify.get_scores(extractor), class_id=num_classes + class_id) print('---------------- Training finished -----------------') with test_utils.temporary_file(suffix='.tflite') as output_model_path: output_model_path.write(engine.serialize_model()) # Evaluate print('---------------- Start evaluating -----------------') classifier = make_interpreter(output_model_path.name) classifier.allocate_tensors() # top[i] represents number of top (i+1) correct inference. top_k_correct_count = [0] * top_k_range image_num = 0 for category, image_list in test_set.items(): n = len(image_list) print('Evaluating {} ({} images)'.format(category, n)) for image_name in image_list: with test_image(os.path.join(dataset_path, category, image_name)) as img: # Set threshold as a negative number to ensure we get top k # candidates even if its score is 0. size = common.input_size(classifier) common.set_input(classifier, img.resize(size, Image.NEAREST)) classifier.invoke() candidates = classify.get_classes(classifier, top_k=top_k_range) for i in range(len(candidates)): candidate = candidates[i] if candidate.id in labels_map and \ labels_map[candidate.id] == category: top_k_correct_count[i] += 1 break image_num += n for i in range(1, top_k_range): top_k_correct_count[i] += top_k_correct_count[i - 1] return [top_k_correct_count[i] / image_num for i in range(top_k_range)]
def main(): args = _parse_args() engine = ImprintingEngine(args.model_path, keep_classes=args.keep_classes) extractor = make_interpreter(engine.serialize_extractor_model(), device=':0') extractor.allocate_tensors() shape = common.input_size(extractor) print('--------------- Parsing data set -----------------') print('Dataset path:', args.data) train_set, test_set = _read_data(args.data, args.test_ratio) print('Image list successfully parsed! Category Num = ', len(train_set)) print('---------------- Processing training data ----------------') print('This process may take more than 30 seconds.') train_input = [] labels_map = {} for class_id, (category, image_list) in enumerate(train_set.items()): print('Processing category:', category) train_input.append( _prepare_images(image_list, os.path.join(args.data, category), shape)) labels_map[class_id] = category print('---------------- Start training -----------------') num_classes = engine.num_classes for class_id, tensors in enumerate(train_input): for tensor in tensors: common.set_input(extractor, tensor) extractor.invoke() embedding = classify.get_scores(extractor) engine.train(embedding, class_id=num_classes + class_id) print('---------------- Training finished! -----------------') with open(args.output, 'wb') as f: f.write(engine.serialize_model()) print('Model saved as : ', args.output) _save_labels(labels_map, args.output) print('------------------ Start evaluating ------------------') interpreter = make_interpreter(args.output) interpreter.allocate_tensors() size = common.input_size(interpreter) top_k = 5 correct = [0] * top_k wrong = [0] * top_k for category, image_list in test_set.items(): print('Evaluating category [', category, ']') for img_name in image_list: img = Image.open(os.path.join(args.data, category, img_name)).resize( size, Image.NEAREST) common.set_input(interpreter, img) interpreter.invoke() candidates = classify.get_classes(interpreter, top_k, score_threshold=0.1) recognized = False for i in range(top_k): if i < len(candidates) and labels_map[ candidates[i].id] == category: recognized = True if recognized: correct[i] = correct[i] + 1 else: wrong[i] = wrong[i] + 1 print('---------------- Evaluation result -----------------') for i in range(top_k): print('Top {} : {:.0%}'.format(i + 1, correct[i] / (correct[i] + wrong[i])))