diff --git a/src/eynollah/model_zoo/model_zoo.py b/src/eynollah/model_zoo/model_zoo.py index d6b5e38..3fca088 100644 --- a/src/eynollah/model_zoo/model_zoo.py +++ b/src/eynollah/model_zoo/model_zoo.py @@ -19,7 +19,7 @@ MODEL_VRAM_LIMITS = { "enhancement": 980, # due to bs 3 "col_classifier": 210, "page": 618, - "textline": 1680, # 954 for bs 1 + "textline": 1880, # 954 for bs 1 "region_1_2": 1580, "region_fl_np": 1756, "table": 1818, @@ -338,6 +338,14 @@ class EynollahModelZoo: # 'cudnn_conv_algo_search': 'EXHAUSTIVE', #'cudnn_conv_use_max_workspace': 0, # 'do_copy_in_default_stream': True, + # enable_cuda_graph + # cudnn_conv1d_pad_to_nc1d + # prefer_nhwc + # tunable_op_enable + # tunable_op_tuning_enable + # tunable_op_max_tuning_duration_ms + # use_ep_level_unified_stream + # enable_skip_layer_norm_strict_mode # ... })] + providers if 'TensorrtExecutionProvider' in providers: @@ -347,9 +355,28 @@ class EynollahModelZoo: 'device_id': gpu, 'trt_max_workspace_size': MODEL_VRAM_LIMITS[model_category] * 1024 * 1024, # 'trt_fp16_enable': True, - # 'trt_engine_cache_enable': True, - # 'trt_timing_cache_enable': True, + # trt_bf16_enable + 'trt_engine_cache_enable': True, + 'trt_timing_cache_enable': True, # ... + # trt_engine_hw_compatible + # trt_engine_cache_path + # trt_engine_cache_prefix + # trt_timing_cache_path + # trt_onnx_model_folder_path + # trt_ep_context_file_path + # trt_cuda_graph_enable + # trt_profile_opt_shapes + # trt_profile_min_shapes + # trt_profile_max_shapes + # trt_builder_optimization_level + # trt_build_heuristics_enable + # trt_sparsity_enable + # trt_weight_stripped_engine_enable + # trt_dla_core + # trt_dla_enable + # trt_min_subgraph_size + # trt_ep_context_embed_mode })] + providers provider0 = providers[0] if isinstance(provider0, tuple): diff --git a/src/eynollah/patch_encoder.py b/src/eynollah/patch_encoder.py index 610f0b4..556a4c9 100644 --- a/src/eynollah/patch_encoder.py +++ b/src/eynollah/patch_encoder.py @@ -29,7 +29,13 @@ class Patches(layers.Layer): self.patch_size_y = patch_size_y def call(self, images): - batch_size = tf.shape(images)[0] + #batch_size = tf.shape(images)[0] + return tf.map_fn(self.call_single, images) + + def call_single(self, image): + # avoid batched extract_patches: too much memory, + # and variable batch dim not supported by ONNX implementation + images = tf.expand_dims(image, axis=0) patches = tf.image.extract_patches( images=images, sizes=[1, self.patch_size_y, self.patch_size_x, 1], @@ -37,8 +43,9 @@ class Patches(layers.Layer): rates=[1, 1, 1, 1], padding="VALID", ) - patch_dims = patches.shape[-1] - return tf.reshape(patches, [batch_size, -1, patch_dims]) + _, n_rows, n_cols, patch_dims = patches.shape + n_tiles = patches.shape[1] * patches.shape[2] #-1 + return tf.reshape(patches, [1, n_tiles, patch_dims]) def get_config(self): return dict(patch_size_x=self.patch_size_x,