A simple wrapper to easily create "linear" graph, consisting of layers / symbolic functions with only one input & output.
| 11 | |
| 12 | |
| 13 | class LinearWrap(object): |
| 14 | """ A simple wrapper to easily create "linear" graph, |
| 15 | consisting of layers / symbolic functions with only one input & output. |
| 16 | """ |
| 17 | |
| 18 | class _TFModuleFunc(object): |
| 19 | def __init__(self, mod, tensor): |
| 20 | self._mod = mod |
| 21 | self._t = tensor |
| 22 | |
| 23 | def __getattr__(self, name): |
| 24 | ret = getattr(self._mod, name) |
| 25 | if isinstance(ret, ModuleType): |
| 26 | return LinearWrap._TFModuleFunc(ret, self._t) |
| 27 | else: |
| 28 | # assume to be a tf function |
| 29 | def f(*args, **kwargs): |
| 30 | o = ret(self._t, *args, **kwargs) |
| 31 | return LinearWrap(o) |
| 32 | return f |
| 33 | |
| 34 | def __init__(self, tensor): |
| 35 | """ |
| 36 | Args: |
| 37 | tensor (tf.Tensor): the tensor to wrap |
| 38 | """ |
| 39 | self._t = tensor |
| 40 | |
| 41 | def __getattr__(self, layer_name): |
| 42 | layer = get_registered_layer(layer_name) |
| 43 | if layer is not None: |
| 44 | # this is a registered tensorpack layer |
| 45 | # parse arguments by tensorpack model convention |
| 46 | if layer.use_scope: |
| 47 | def layer_func(name, *args, **kwargs): |
| 48 | ret = layer(name, self._t, *args, **kwargs) |
| 49 | return LinearWrap(ret) |
| 50 | else: |
| 51 | def layer_func(*args, **kwargs): |
| 52 | if len(args) and isinstance(args[0], six.string_types): |
| 53 | name, args = args[0], args[1:] |
| 54 | ret = layer(name, self._t, *args, **kwargs) |
| 55 | else: |
| 56 | ret = layer(self._t, *args, **kwargs) |
| 57 | return LinearWrap(ret) |
| 58 | return layer_func |
| 59 | else: |
| 60 | assert layer_name == 'tf', \ |
| 61 | "Calling LinearWrap.{}:" \ |
| 62 | " neither a layer nor 'tf'! " \ |
| 63 | "Did you forget to extract tensor from LinearWrap?".format(layer_name) |
| 64 | import tensorflow as layer # noqa |
| 65 | assert isinstance(layer, ModuleType), layer |
| 66 | return LinearWrap._TFModuleFunc(layer, self._t) |
| 67 | |
| 68 | def apply(self, func, *args, **kwargs): |
| 69 | """ |
| 70 | Apply a function on the wrapped tensor. |
no outgoing calls
no test coverage detected