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
|
"enhancement": 980, # due to bs 3
|
||||||
"col_classifier": 210,
|
"col_classifier": 210,
|
||||||
"page": 618,
|
"page": 618,
|
||||||
"textline": 1680, # 954 for bs 1
|
"textline": 1880, # 954 for bs 1
|
||||||
"region_1_2": 1580,
|
"region_1_2": 1580,
|
||||||
"region_fl_np": 1756,
|
"region_fl_np": 1756,
|
||||||
"table": 1818,
|
"table": 1818,
|
||||||
|
|
@ -338,6 +338,14 @@ class EynollahModelZoo:
|
||||||
# 'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
# 'cudnn_conv_algo_search': 'EXHAUSTIVE',
|
||||||
#'cudnn_conv_use_max_workspace': 0,
|
#'cudnn_conv_use_max_workspace': 0,
|
||||||
# 'do_copy_in_default_stream': True,
|
# '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
|
})] + providers
|
||||||
if 'TensorrtExecutionProvider' in providers:
|
if 'TensorrtExecutionProvider' in providers:
|
||||||
|
|
@ -347,9 +355,28 @@ class EynollahModelZoo:
|
||||||
'device_id': gpu,
|
'device_id': gpu,
|
||||||
'trt_max_workspace_size': MODEL_VRAM_LIMITS[model_category] * 1024 * 1024,
|
'trt_max_workspace_size': MODEL_VRAM_LIMITS[model_category] * 1024 * 1024,
|
||||||
# 'trt_fp16_enable': True,
|
# 'trt_fp16_enable': True,
|
||||||
# 'trt_engine_cache_enable': True,
|
# trt_bf16_enable
|
||||||
# 'trt_timing_cache_enable': True,
|
'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
|
})] + providers
|
||||||
provider0 = providers[0]
|
provider0 = providers[0]
|
||||||
if isinstance(provider0, tuple):
|
if isinstance(provider0, tuple):
|
||||||
|
|
|
||||||
|
|
@ -29,7 +29,13 @@ class Patches(layers.Layer):
|
||||||
self.patch_size_y = patch_size_y
|
self.patch_size_y = patch_size_y
|
||||||
|
|
||||||
def call(self, images):
|
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(
|
patches = tf.image.extract_patches(
|
||||||
images=images,
|
images=images,
|
||||||
sizes=[1, self.patch_size_y, self.patch_size_x, 1],
|
sizes=[1, self.patch_size_y, self.patch_size_x, 1],
|
||||||
|
|
@ -37,8 +43,9 @@ class Patches(layers.Layer):
|
||||||
rates=[1, 1, 1, 1],
|
rates=[1, 1, 1, 1],
|
||||||
padding="VALID",
|
padding="VALID",
|
||||||
)
|
)
|
||||||
patch_dims = patches.shape[-1]
|
_, n_rows, n_cols, patch_dims = patches.shape
|
||||||
return tf.reshape(patches, [batch_size, -1, patch_dims])
|
n_tiles = patches.shape[1] * patches.shape[2] #-1
|
||||||
|
return tf.reshape(patches, [1, n_tiles, patch_dims])
|
||||||
|
|
||||||
def get_config(self):
|
def get_config(self):
|
||||||
return dict(patch_size_x=self.patch_size_x,
|
return dict(patch_size_x=self.patch_size_x,
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue