def models(): tf.keras.backend.set_learning_phase(False) tf_test_utils.set_random_seed() input_shape = get_input_shape(FLAGS.data) # keras model receives images size as input, # where batch size is not specified - by default it is dynamic if FLAGS.model in APP_MODELS: weights = 'imagenet' if FLAGS.data == 'imagenet' else None # TODO(rybakov) with the fix of https://github.com/google/iree/issues/1660 # add include_top=True # if weights == 'imagenet' it will load weights from external tf.keras URL model = APP_MODELS[FLAGS.model](weights=weights, input_shape=input_shape[1:]) if FLAGS.data == 'cifar10' and FLAGS.url: file_name = 'cifar10' + FLAGS.model # it will download model weights from publically available folder: PATH # and save it to cache_dir=~/.keras and return path to it weights_path = tf.keras.utils.get_file( file_name, os.path.join(FLAGS.url, file_name + '.h5')) model.load_weights(weights_path) else: raise ValueError('Unsupported model', FLAGS.model) module = tf.Module() module.m = model # specify input size with static batch size # TODO(b/142948097): with support of dynamic shape # replace input_shape by model.input_shape, so batch size will be dynamic (-1) module.predict = tf.function(input_signature=[tf.TensorSpec(input_shape)])( model.call) return module
def __init__(self, build_fn, *args, also_track=None, **kwargs): """Defers initialization of an object with transformed arguments. Args: build_fn: Python callable specifying a deferred transformation of the provided arguments. This must have signature `module = build_fn(*args, **kwargs)`. The return value `module` is an instance of `tf.Module`. *args: Optional positional arguments to `build_fn`. also_track: Optional instance or structure of instances of `tf.Variable` and/or `tf.Module`, containing any additional trainable variables that the `build_fn` may access beyond the given `args` and `kwargs`. This ensures that such variables will be correctly tracked in `self.trainable_variables`. Default value: `None`. **kwargs: Optional keyword arguments to `build_fn`. """ self._build_fn = build_fn self._param_args = args self._param_kwargs = kwargs self._deferred_module_also_track = also_track # In order for DeferredModule to work as a tf.Module, we need to ensure that # attrs used by tf.Module are handled directly, rather than being forwarded # to the inner class. self._module_attrs = set(dir(tf.Module())) super(DeferredModule, self).__init__()
def __init__(self, base_class, args_fn, *args, **kwargs): """Defers initialization of an object with transformed arguments. Args: base_class: Python type or callable such that `base_class(**args_fn(...))` is an instance of `tf.Module`---for example, a TFP Distribution or Bijector. args_fn: Python callable specifying a deferred transformation of the provided arguments. This must have signature `base_class_init_args = args_fn(*args, **kwargs)`. The return value `base_class_init_args` may be either a dictionary or an iterable (list/tuple), in which case the class will be initialized as `base_class(**base_class_init_args)` or `base_class(*base_class_init_args)`, respectively. *args: Optional positional arguments to `args_fn`. **kwargs: Optional keyword arguments to `args_fn`. """ self._base_class = base_class self._args_fn = args_fn self._param_args = args self._param_kwargs = kwargs # In order for DeferredModule to work as a tf.Module, we need to ensure that # attrs used by tf.Module are handled directly, rather than being forwarded # to the inner class. self._module_attrs = set(dir(tf.Module())) super(DeferredModule, self).__init__()
def _testConvertedFunction(self, obj, func, converted_concrete_func, input_data): # Ensure the converted graph has no variables and no function calls. constant_graph_def = converted_concrete_func.graph.as_graph_def() self.assertEqual(0, self._getNumVariables(constant_graph_def)) self.assertFalse( self._hasStatefulPartitionedCallOp(constant_graph_def)) # Check that the converted ConcreteFunction produces the same result as the # original Function. expected_value = tf.nest.flatten(func(**input_data)) actual_value = tf.nest.flatten(converted_concrete_func(**input_data)) for expected, actual in zip(expected_value, actual_value): np.testing.assert_almost_equal(expected.numpy(), actual.numpy()) # Ensure the shape is retained. for tensor in converted_concrete_func.inputs: actual_shape = input_data[tensor.name.split(":")[0]].shape self.assertEqual(tensor.shape, actual_shape) # Save the converted ConcreteFunction as a signature. save_dir = os.path.join(self.get_temp_dir(), "frozen_saved_model") root = tf.Module() root.f = converted_concrete_func save(root, save_dir, {"mykey": converted_concrete_func}) # Load it back and make sure it works. loaded_obj = load(save_dir) actual_value = tf.nest.flatten( loaded_obj.signatures["mykey"](**input_data)) for expected, actual in zip(expected_value, actual_value): np.testing.assert_almost_equal(expected.numpy(), actual.numpy())
def test_model_wrapped_in_module_discovers_submodules(self): linear = tf.keras.models.Sequential( [tf.keras.layers.Dense(units=1, input_shape=[1])]) linear.compile(optimizer="sgd", loss="mean_squared_error") m = tf.Module() m.l = linear self.assertNotEmpty(m.submodules) self.assertLen(m.variables, 2)
def testDictWrapperBadKeys(self): a = tf.Module() a.d = {} a.d[1] = data_structures.wrap_or_unwrap([]) model = training.Model() model.sub = a save_path = os.path.join(self.get_temp_dir(), "ckpt") with self.assertRaisesRegex(ValueError, "non-string key"): model.save_weights(save_path)
def testDictWrapperNoDependency(self): a = tf.Module() a.d = data_structures.NoDependency({}) a.d[1] = [3] self.assertEqual([a], util.list_objects(a)) model = training.Model() model.sub = a save_path = os.path.join(self.get_temp_dir(), "ckpt") model.save_weights(save_path) model.load_weights(save_path)
def testNoDependency(self): root = tf.Module() hasdep = tf.Module() root.hasdep = hasdep nodep = tf.Module() root.nodep = data_structures.NoDependency(nodep) self.assertEqual(1, len(root._checkpoint_dependencies)) self.assertIs(root._checkpoint_dependencies[0].ref, root.hasdep) self.assertIs(root.hasdep, hasdep) self.assertIs(root.nodep, nodep) class NoDependencyModel(training.Model): @base.no_automatic_dependency_tracking def __init__(self): super(NoDependencyModel, self).__init__() self.a = [] self.b = tf.Module() nodeps = NoDependencyModel() self.assertEqual([nodeps], util.list_objects(nodeps))
def testNoDependency(self): root = tf.Module() hasdep = tf.Module() root.hasdep = hasdep nodep = tf.Module() root.nodep = data_structures.NoDependency(nodep) self.assertLen(root._trackable_children(), 1) self.assertIs(root._trackable_children()["hasdep"], root.hasdep) self.assertIs(root.hasdep, hasdep) self.assertIs(root.nodep, nodep) class NoDependencyModel(training.Model): @tf.__internal__.tracking.no_automatic_dependency_tracking def __init__(self): super(NoDependencyModel, self).__init__() self.a = [] self.b = tf.Module() nodeps = NoDependencyModel() self.assertEqual([nodeps], util.list_objects(nodeps))
def testNoDepList(self): a = training.Model() a.l1 = data_structures.NoDependency([]) a.l1.insert(1, 0) self.assertIsInstance(a.l1, list) checkpoint = tf.train.Checkpoint(a=a) checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt")) a.l2 = [] a.l2.insert(1, tf.Module()) with self.assertRaisesRegex(ValueError, "A list element was replaced"): checkpoint.save(os.path.join(self.get_temp_dir(), "ckpt"))
def testNonAppendNotTrackable(self): # Non-append mutations (deleting or overwriting values) are OK when the # values aren't tracked. a = tf.Module() a.d = {} a.d["a"] = [3] a.d[1] = 3 a.d[1] = 2 self.assertEqual(2, a.d[1]) del a.d[1] a.d[2] = data_structures.NoDependency(tf.Module()) second = tf.Module() a.d[2] = data_structures.NoDependency(second) self.assertIs(second, a.d[2]) self.assertEqual([a, a.d, a.d["a"]], util.list_objects(a)) model = training.Model() model.sub = a save_path = os.path.join(self.get_temp_dir(), "ckpt") model.save_weights(save_path) model.load_weights(save_path)
def testNonStringKeyNotTrackableValue(self): a = tf.Module() a.d = {} a.d["a"] = [3] a.d[1] = data_structures.NoDependency([3]) self.assertEqual([a, a.d, a.d["a"]], util.list_objects(a)) model = training.Model() model.sub = a save_path = os.path.join(self.get_temp_dir(), "ckpt") model.save_weights(save_path) model.load_weights(save_path)
def lstm_module(): tf_utils.set_random_seed() inputs = tf.keras.layers.Input(batch_size=NUM_BATCH, shape=INPUT_SHAPE[1:]) outputs = tf.keras.layers.LSTM(units=NUM_UNITS, return_sequences=True)(inputs) model = tf.keras.Model(inputs, outputs) module = tf.Module() module.m = model module.predict = tf.function( input_signature=[tf.TensorSpec(INPUT_SHAPE, tf.float32)])( model.call) return module
def testLayerCollectionWithExternalMutation(self): d = {} root = tf.Module() root.wrapper = d self.assertEqual([], root.wrapper.layers) self.assertEqual([], root.wrapper.trainable_weights) layer1 = core.Dense(1) layer2 = core.Dense(1) d["a"] = layer1 d["b"] = layer2 self.assertEqual([layer1, layer2], root.wrapper.layers) # The layers have still not created variables self.assertEqual([], root.wrapper.trainable_weights)
def _freezeModel(self, model): """Freezes the model. Args: model: Function. Returns: root: AutoTrackable object with original ConcreteFunction. output_func: frozen ConcreteFunction. """ root = tf.Module() root.f = model input_func = root.f.get_concrete_function() output_func = convert_to_constants.convert_variables_to_constants_v2( input_func, lower_control_flow=False) return root, output_func
def test_module_discover_layer_variable(self): m = tf.Module() m.a = tf.keras.layers.Dense(1) m.b = tf.keras.layers.Dense(2) # The weights of the layer has not been created yet. self.assertEmpty(m.variables) self.assertLen(m.submodules, 2) inputs = tf.keras.layers.Input((1, )) m.a(inputs) m.b(inputs) variable_list = m.variables self.assertLen(variable_list, 4) self.assertIs(variable_list[0], m.a.kernel) self.assertIs(variable_list[1], m.a.bias) self.assertIs(variable_list[2], m.b.kernel) self.assertIs(variable_list[3], m.b.bias)
def models(): tf.keras.backend.set_learning_phase(False) # keras model receives images size as input, # where batch size is not specified - by default it is dynamic if FLAGS.model in APP_MODELS: model = APP_MODELS[FLAGS.model](weights=None, include_top=False, input_shape=INPUT_SHAPE[1:]) else: raise ValueError('unsupported model', FLAGS.model) module = tf.Module() module.m = model # specify input size with static batch size # TODO(b/142948097): with support of dynamic shape # replace INPUT_SHAPE by model.input_shape, so batch size will be dynamic (-1) module.predict = tf.function(input_signature=[tf.TensorSpec(INPUT_SHAPE)])( model.call) return module
def __init__(self, build_fn, *args, **kwargs): """Defers initialization of an object with transformed arguments. Args: build_fn: Python callable specifying a deferred transformation of the provided arguments. This must have signature `module = build_fn(*args, **kwargs)`. The return value `module` is an instance of `tf.Module`. *args: Optional positional arguments to `build_fn`. **kwargs: Optional keyword arguments to `build_fn`. """ self._build_fn = build_fn self._param_args = args self._param_kwargs = kwargs # In order for DeferredModule to work as a tf.Module, we need to ensure that # attrs used by tf.Module are handled directly, rather than being forwarded # to the inner class. self._module_attrs = set(dir(tf.Module())) super(DeferredModule, self).__init__()
def __init__(self): super().__init__() self.a = [] self.b = tf.Module()
def __init__(self): super(NoDependencyModel, self).__init__() self.a = [] self.b = tf.Module()