"...git@developer.sourcefind.cn:renzhc/diffusers_dcu.git" did not exist on "33e636cea557d73d0ff4387c3096375ea00ab0d4"
Commit 4f58b1f7 authored by A. Unique TensorFlower's avatar A. Unique TensorFlower
Browse files

Internal change

PiperOrigin-RevId: 301434099
parent bc5b60ef
...@@ -41,8 +41,14 @@ class TfHubMemoryUsageBenchmark(PerfZeroBenchmark): ...@@ -41,8 +41,14 @@ class TfHubMemoryUsageBenchmark(PerfZeroBenchmark):
output_dir=output_dir, default_flags=default_flags, **kwargs) output_dir=output_dir, default_flags=default_flags, **kwargs)
if hub_model_handle_list: if hub_model_handle_list:
for hub_model_handle in hub_model_handle_list.split(';'): for hub_model_handle in hub_model_handle_list.split(';'):
# Converts a model handle of the form
# https://tfhub.dev/google/nnlm-en-dim128/1 to valid python method name
# like google_nnlm_en_dim128_1.
hub_model_method_name = hub_model_handle.replace(
'https://tfhub.dev',
'').replace('/', '_').replace('-', '_').strip('_')
setattr( setattr(
self, 'benchmark_' + hub_model_handle, self, 'benchmark_' + hub_model_method_name,
functools.partial(self.benchmark_memory_usage, hub_model_handle)) functools.partial(self.benchmark_memory_usage, hub_model_handle))
def benchmark_memory_usage( def benchmark_memory_usage(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment