Spaces:
Runtime error
Runtime error
| # Copyright 2023 The TensorFlow Authors. All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| """A script to export a TF-Hub SavedModel.""" | |
| from typing import List, Optional | |
| # Import libraries | |
| import tensorflow as tf, tf_keras | |
| from official.core import config_definitions as cfg | |
| from official.vision import configs | |
| from official.vision.modeling import factory | |
| def build_model(batch_size: Optional[int], | |
| input_image_size: List[int], | |
| params: cfg.ExperimentConfig, | |
| num_channels: int = 3, | |
| skip_logits_layer: bool = False) -> tf_keras.Model: | |
| """Builds a model for TF Hub export. | |
| Args: | |
| batch_size: The batch size of input. | |
| input_image_size: A list of [height, width] specifying the input image size. | |
| params: The config used to train the model. | |
| num_channels: The number of input image channels. | |
| skip_logits_layer: Whether to skip the logits layer for image classification | |
| model. Default is False. | |
| Returns: | |
| A tf_keras.Model instance. | |
| Raises: | |
| ValueError: If the task is not supported. | |
| """ | |
| input_specs = tf_keras.layers.InputSpec(shape=[batch_size] + | |
| input_image_size + [num_channels]) | |
| if isinstance(params.task, | |
| configs.image_classification.ImageClassificationTask): | |
| model = factory.build_classification_model( | |
| input_specs=input_specs, | |
| model_config=params.task.model, | |
| l2_regularizer=None, | |
| skip_logits_layer=skip_logits_layer) | |
| else: | |
| raise ValueError('Export module not implemented for {} task.'.format( | |
| type(params.task))) | |
| return model | |
| def export_model_to_tfhub(batch_size: Optional[int], | |
| input_image_size: List[int], | |
| params: cfg.ExperimentConfig, | |
| checkpoint_path: str, | |
| export_path: str, | |
| num_channels: int = 3, | |
| skip_logits_layer: bool = False): | |
| """Export a TF2 model to TF-Hub.""" | |
| model = build_model(batch_size, input_image_size, params, num_channels, | |
| skip_logits_layer) | |
| checkpoint = tf.train.Checkpoint(model=model) | |
| checkpoint.restore(checkpoint_path).assert_existing_objects_matched() | |
| model.save(export_path, include_optimizer=False, save_format='tf') | |