mirror of
https://github.com/qurator-spk/eynollah.git
synced 2026-07-11 22:29:29 +02:00
models (ViT backbone) iterate extract_patches over batch dim…
`Patches.call`: use `tf.map_fn` instead of running entire batch through `tf.image.extract_patches` (faster, less VRAM, allows ONNX conversion to work)
This commit is contained in:
parent
1b27c7390f
commit
16943f70b4
2 changed files with 40 additions and 6 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue