Spaces:
Runtime error
Runtime error
File size: 3,504 Bytes
5672777 93528c6 5672777 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 |
# 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.
"""Backbone registers and factory method.
One can regitered a new backbone model by the following two steps:
1 Import the factory and register the build in the backbone file.
2 Import the backbone class and add a build in __init__.py.
```
# my_backbone.py
from modeling.backbones import factory
class MyBackbone():
...
@factory.register_backbone_builder('my_backbone')
def build_my_backbone():
return MyBackbone()
# backbones/__init__.py adds import
from modeling.backbones.my_backbone import MyBackbone
```
If one wants the MyBackbone class to be used only by those binary
then don't imported the backbone module in backbones/__init__.py, but import it
in place that uses it.
"""
from typing import Sequence, Union
# Import libraries
import tensorflow as tf, tf_keras
from official.core import registry
from official.modeling import hyperparams
_REGISTERED_BACKBONE_CLS = {}
def register_backbone_builder(key: str):
"""Decorates a builder of backbone class.
The builder should be a Callable (a class or a function).
This decorator supports registration of backbone builder as follows:
```
class MyBackbone(tf_keras.Model):
pass
@register_backbone_builder('mybackbone')
def builder(input_specs, config, l2_reg):
return MyBackbone(...)
# Builds a MyBackbone object.
my_backbone = build_backbone_3d(input_specs, config, l2_reg)
```
Args:
key: A `str` of key to look up the builder.
Returns:
A callable for using as class decorator that registers the decorated class
for creation from an instance of task_config_cls.
"""
return registry.register(_REGISTERED_BACKBONE_CLS, key)
def build_backbone(input_specs: Union[tf_keras.layers.InputSpec,
Sequence[tf_keras.layers.InputSpec]],
backbone_config: hyperparams.Config,
norm_activation_config: hyperparams.Config,
l2_regularizer: tf_keras.regularizers.Regularizer = None,
**kwargs) -> tf_keras.Model: # pytype: disable=annotation-type-mismatch # typed-keras
"""Builds backbone from a config.
Args:
input_specs: A (sequence of) `tf_keras.layers.InputSpec` of input.
backbone_config: A `OneOfConfig` of backbone config.
norm_activation_config: A config for normalization/activation layer.
l2_regularizer: A `tf_keras.regularizers.Regularizer` object. Default to
None.
**kwargs: Additional keyword args to be passed to backbone builder.
Returns:
A `tf_keras.Model` instance of the backbone.
"""
backbone_builder = registry.lookup(_REGISTERED_BACKBONE_CLS,
backbone_config.type)
return backbone_builder(
input_specs=input_specs,
backbone_config=backbone_config,
norm_activation_config=norm_activation_config,
l2_regularizer=l2_regularizer,
**kwargs)
|