Discriminative Layer Training #969

hyang0129 · 2020-01-29T18:10:42Z

See #958

I will continue to add more tests and setup distributed tests as well.

If you have any feedback, please let me know

hyang0129 · 2020-01-29T18:12:02Z

I think I forgot to run flake8 and black

hyang0129 · 2020-01-29T18:13:18Z

@googlebot I signed it!

googlebot · 2020-01-29T18:13:27Z

CLAs look good, thanks!

ℹ️ Googlers: Go here for more info.

…ild and init

hyang0129 · 2020-09-09T19:46:00Z

@WindQAQ please review

WindQAQ

Thanks for the contribution! Generally LGTM. Some opinions here:

We should elaborate doctoring that layers should be built before constructing this optimizer.
Instead of holding name of variables, we can use .ref(). https://www.tensorflow.org/api_docs/python/tf/Variable#ref
Can we override __repr__ to give more information when printing it?

WindQAQ · 2020-09-10T02:51:32Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+from tensorflow.keras.optimizers import Optimizer
+
+# python -m flake8 tensorflow_addons/optimizers/discriminative_layer_training.py
+# python -m black tensorflow_addons/optimizers/discriminative_layer_training.py


Remove these.

WindQAQ · 2020-09-10T02:51:34Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+class FakeVar:
+    def __init__(self, name):
+        # probably can be refactored out
+        self.name = name


any reason that we should use this?

I'm open to alternative ideas. Its just an object to match how we access the name of weight object

Shouldn't tf.Variable itself have name attribute?

I wanted to avoid serializing the whole variable as part of the optimizer, so I used the name. This has been refactored to remove fakevar. I tried using var.ref but it dies on deserialization

WindQAQ · 2020-09-10T02:53:05Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+
+
+@tf.keras.utils.register_keras_serializable(package="Addons")
+class MultiOpt(Optimizer):


Can we have a better name? At least call it MultiOptimizer instead of abbreviation.

WindQAQ · 2020-09-10T02:54:15Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+            )
+
+        self.initialized_optimizer_specs = [
+            self.initialize_from_optimizer_spec(spec) for spec in self.optimizer_specs


Why we need to do this get_config -> from_config?

to prevent serialization issues when running on TPU. I don't want to pass instantiated optimizers for graph construction on the TPU. This should protect from weird problems on serialization and deserialization. As in, when we need to initialize the multiopt from a serialized state, we cannot pass instantiated optimizers to the multiopt

Thanks for the clarification 😃 It makes sense to me! Can we also have this in docstring? Thank you!

so I rewrote some stuff that line is no longer needed after fixing the serialization

WindQAQ · 2020-09-10T02:54:56Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+        for spec in self.optimizer_specs:
+            spec["gv"] = []
+
+        for grad, var in tuple(grads_and_vars):


tuple is unnecessary right?

I believe it matches the parent class's usage of tuple

WindQAQ · 2020-09-10T02:56:37Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+                FakeVar(var.name) for sublayer in layer for var in sublayer.weights
+            ]
+        else:
+            weights = [FakeVar(var.name) for var in layer.weights]


trainable_weights maybe?

I'm not sure if that would make a difference because the wrapper calls the apply function of the wrapped optimizers. I could try it, but I believe that the wrapped optimizers have their own way of handling trainable vs non trainable weights in their apply step. I'm guessing that a normal opt receives [weights] when running model.fit so either the model.fit only passes trainable weights or the opt.apply filters them out somehow. If that's the case, then it would be unnecessary to filter the weights twice.

The optimizer will not distinguish if the variable is trainable or not. It will create slots and compute gradient w.r.t. all variables passing to apply_gradients.

https://colab.research.google.com/drive/1NOigO7WuGdFKTKD4XovqD52SMrSYzI99?usp=sharing

I think maybe that's intended behavior when using the optimizer outside of the model.fit loop? When used in the model.fit loop, the multi opt correctly ignores layers without the trainable attribute.

I would prefer if the multi optimizer acts solely as a wrapper and also passes forward all bugs/features of the wrapped optimizer.

link demonstrating layer trainable false works for multi optimizer here

Got it. It looks good to me if that won't cause any errors.

WindQAQ · 2020-09-10T02:56:39Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+
+        return {
+            "optimizer_class": optimizer_instance.__class__,
+            "optimizer_config": optimizer_instance.get_config(),


we can use tf.keras.optimizers.serialize and tf.keras.optimizers.deserialize.

yeah I found a less bad way to organize the wrapper. will push when testing is done

WindQAQ · 2020-09-10T02:56:44Z

tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py

@@ -0,0 +1,101 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.


tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py

WindQAQ · 2020-09-10T02:58:21Z

tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py

+    if system == "Windows":
+        return [tf.float32, tf.float64]
+    else:
+        return _dtypes_to_test(use_gpu)


do we have tf.linalg.svd in our test?

I don't think so? I'm not sure what that is TBH. I copied some code from another optimizer test so I can remove this.

Sure, just remove it.

WindQAQ · 2020-09-10T03:29:12Z

Also, I would like to have some tests on saving this optimizer

https://www.tensorflow.org/api_docs/python/tf/train/Checkpoint

Probably we should use self._track_trackable(optimizer, name="xxx") to keep track on those optimizers.

hyang0129 · 2020-09-13T15:47:38Z

@WindQAQ I'll look into that for optimizer saving.

hyang0129 · 2020-09-13T23:40:28Z

@WindQAQ what kind of info would be helpful when calling print(multioptimizer)?

WindQAQ · 2020-09-13T23:48:01Z

@WindQAQ what kind of info would be helpful when calling print(multioptimizer)?

I guess users may want to see the optimizers and their corresponding variables. This is optional, just depend on your choice if you also think it is useful.

hyang0129 · 2020-09-14T00:18:29Z

@WindQAQ should be good for review

WindQAQ

Generally LGTM! Some minor documentation and naming issue :-) Thank you again for bringing this up!

WindQAQ · 2020-09-14T00:28:05Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+    @typechecked
+    def __init__(
+        self,
+        optimizer_layer_pairs: Union[list, None] = None,


I think optimizers_and_layers much fit the tensorflow naming philosophy like apply_gradients(grads_and_vars).

WindQAQ · 2020-09-14T00:31:59Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+        **kwargs
+    ):
+
+        """


r"""Short descriptions in the first line, do not break line so rendering can work correctly. Detailed descriptions start here. Args: foo: foo bar: bar Usage: >>> code snippet here >>> code snippet here Note: Note here if any. Reference: - [Name of the paper](linktothepaper) """

WindQAQ · 2020-09-14T00:32:07Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+        ```
+
+
+        """


This should be put between class MultiOptimizer and def __init__. We also miss a section of Args. It's also encouraged to use doctest for python code snippet (optional, if you have time)

See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/keras/optimizer_v2/adam.py#L47-L69

WindQAQ · 2020-09-14T00:33:36Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+import tensorflow as tf
+from typeguard import typechecked
+from typing import Union
+from tensorflow.keras.optimizers import Optimizer


from typing import Union import tensorflow as tf from tensorflow.keras.optimizers import Optimizer from typeguard import typechecked

WindQAQ · 2020-09-14T00:34:11Z

tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py

+
+from tensorflow_addons.optimizers.discriminative_layer_training import MultiOptimzer
+import tensorflow as tf
+from tensorflow_addons.utils import test_utils


import pytest import numpy as np import tensorflow as tf from tensorflow_addons.optimizers.discriminative_layer_training import MultiOptimzer from tensorflow_addons.utils import test_utils

WindQAQ · 2020-09-14T00:35:09Z

tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py

+    # The function "_DtypesToTest" is from
+    #       "https://github.com/tensorflow/tensorflow/blob/5d4a6cee737a1dc6c20172a1dc1
+    #        5df10def2df72/tensorflow/python/kernel_tests/conv_ops_3d_test.py#L53-L62"
+    # TODO(WindQAQ): xxx


sorry for misleading. can we change xxx to Clean up this in TF2.4? Thank you!

WindQAQ · 2020-09-14T00:38:16Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+
+
+@tf.keras.utils.register_keras_serializable(package="Addons")
+class MultiOptimzer(Optimizer):


maybe change this to tf.keras.optimizers.Optimizer since it only uses once.

WindQAQ · 2020-09-14T00:44:50Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+
+    @classmethod
+    def create_optimizer_spec(cls, optimizer_instance, layer):
+


Should we check isinstance(optimizer, tf.keras.optimziers.Optimizer) here?

WindQAQ · 2020-09-14T00:52:38Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+        Wrapped Gradient Apply method. Returns a list of tf ops to be executed.
+
+        Name of variable is used rather than var.ref() to enable serialization and deserialization.
+        """


"""Wrapped apply_gradient method. Returns a list of tf ops to be executed. Name of variable is used rather than var.ref() to enable serialization and deserialization. """

WindQAQ · 2020-09-14T00:53:56Z

tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py

+    # expect weights to be same for layer 2
+    test_utils.assert_allclose_according_to_type(
+        weights_before_train[1], weights_after_train[1]
+    )


We should also have a test about serialization.

https://github.com/tensorflow/addons/blob/master/tensorflow_addons/optimizers/tests/lookahead_test.py#L184-L188

I added a serialization to the existing test as a parameter.

After compiling but before fitting, the test serializes the model with the optimizer, clears the session, then loads the model.

nvm I added a serialization test by itself as well

WindQAQ · 2020-09-14T01:12:21Z

tensorflow_addons/optimizers/discriminative_layer_training.py

+                "You must specify either an list of optimizer_layer_pairs or a list of optimizer_specs"
+            )
+
+    def apply_gradients(self, grads_and_vars, name=None):


This should be

def apply_gradients(self, grads_and_vars, **kwargs): ... return tf.group( [ spec["optimizer"].apply_gradients(spec["gv"], **kwargs) for spec in self.optimizer_specs ] )

to prevent upstream change.

… for only the serialization

hyang0129 · 2020-09-14T02:05:35Z

@WindQAQ I've made the changes

WindQAQ

LGTM! Sorry for the very long delay and thanks for the contribution to this good feature! Thank you!

stale

hyang0129 · 2020-09-14T17:39:56Z

LGTM! Sorry for the very long delay and thanks for the contribution to this good feature! Thank you!

I'm glad I could help! Thanks for the feedback to make this possible.

WindQAQ · 2020-09-14T17:42:33Z

@hyang0129 Oops, We forgot to update code owners, would you mind filing another PR to update https://github.com/tensorflow/addons/blob/master/.github/CODEOWNERS ?

WindQAQ · 2020-09-14T17:43:40Z

@hyang0129 Another thing is that MultiOptimzer -> MultiOptimizer.

hyang0129 · 2020-09-14T17:45:04Z

@hyang0129 Another thing is that MultiOptimzer -> MultiOptimizer.

I'm not sure what you mean by that? Could you elaborate?

WindQAQ · 2020-09-14T17:46:36Z

@hyang0129 Another thing is that MultiOptimzer -> MultiOptimizer.

I'm not sure what you mean by that? Could you elaborate?

We misspell the word "Optimizer". We should correct it in next PR.

hyang0129 · 2020-09-14T17:47:57Z

I'll fix that

seanpmorgan · 2020-09-16T04:33:22Z

@hyang0129 Thanks so much for your great contribution! Happy to see it merged, apologies for the delay.

* initial setup. need to build tests * build some tests. need to test them * fixed typo * created first test * created first test * accidentally messed up another file * accidentally messed up another file * accidentally messed up another file * added run all distributed * fixed formatting * trying to fix tests not running on github CI. * realized that I should probably add the new optimizer files to the build and init * added typeguard and docstring * removed run_all_distributed * graph and eager testing for SGD * reformatted * added distributed tests * removed distributed tests * reverted discriminative layer grad adjust back to apply gradients * added distributed tests with one time virtual device init * increased tolerance for distributed added comments explaining tests * changed how distributed is recognized for increasing tolerance * Redesigned Logic into Optimizer Wrapper (tensorflow#1) * redesigned methodology to use multiple optimizers (one per unique LR) and pass grads to these multiple optimizers. Should allow for complex optimizers to behave properly * adjusted behavior of resource apply to only return the op if the lr_mult matches the lr_mult of the optimizer should only return 1 op for each var. * updated init file changed training config * removed variable position and added some more comments * removed grouped variables as unnecessary * reformatted * updated documentation explicitly defined serialization as not supported * added typecheck for name * added typecheck for name * fixed blank line at end of init file * realized no new line meant to add new line guessing that build file needs to be in alpha order? * ran buildifier * fixed accidentally affecting moving average * changed print to logging.info * changed print to logging.info * Revert "changed print to logging.info" This reverts commit 3fa5e19 * added tutorial. tutorial doesn't import from tfa. May need to remove from PR. Please let me know * refactored to use static method refactored to use getattr updated warning on not using lr_mult expanded on some docstrings * updated the usage of lr_mult in variables * renamed discriminative wrapper to disclayeropt * added note to disuade directly calling apply_gradients * updated toy_cnn to use tempdir and no longer call context.eager implemented toy_rnn function with same flow as toycnn * added toy_rnn and sgd to the test permutations * refactored permutes and train results into private fns * reformatted files and fixed flake 8 issues fixed bad references when lr_mult was changed * added missing functions in prep for tests * updated assign lr mult and explained further why refactored get lowest layers to assign sublayers explained recursively assign sublayers better * forgot to run black so ran it to reformat * specified inputshape for rnn * increased size of test temporarily removed SGD opt. Double opts doubles the number of tests to run so just need to see how long this one takes. * remove toy rnn for now * changed back to medium. maybe large was not actually increasing runtime * fixed input layer * fixed input layer being in wrong place * virtual device modification issue * fixed incorrect usage of lr_mult * added comments for tests explaining them better added toy rnn for testing * added new test fix toy rnn initialization * fixed typo * added inputshape so that pretrained rnn generates weights * changed test to allow head to learn. it should move the loss better * reformatted * fixed test for variable assignment added get config and from config * reformatted * fixed layer references from 1 to 0 because input layer isn't counted as an actual layer in the layer list * reformatted * increased lr and epochs because learning was happning, but assertless tolerance too low * attempting to use run distributed from test utils * removed tutorial * switched to alternative distributed training method * trying to use run distributed without graph and eager * trying to use run_distributed * seems that doing any tensorstuff before tf.test.main creates the issue. changed models to auto check if weights exist and create or load * forgot to return a model on first run of model fn * create model weights on init * changed how args are passed for testcase * changed how args are passed for testcase * try fix init * trying to init weights on model properly * trying to init weights on model properly * just trying all the possibilities * trying to fix weights setup * expanded some comments for some tests * fixed some docstrings and expanded on some comments * reformatted files expanded on many comments and added full stops fixed get/from_config based on optimzierv2 added model checkpoint test * capitalized comments properly. * removed sgd, reduced size of training inputs. * simplified checkpoint name * reformatted * remove run tests in notebook * updated README.md fixed indent for __init__ added test for from config and to config * fixed formatting * removed distributed tests and added a warning if optimizer is initialized within a strategy scope * renamed test_wrap to wrap_test bc pytest thought it was a test. * converting tests into the pytest framework * converted tests and parameterized * cleaned up code * added additional checks and doc string for changes in lr multiplier during training. * changed comment * Simplified discriminative layer training by using a multi optimizer wrapper class. Removed old tests and added new tests conforming to pytest standard. * Refactored code using black and flake8 * updated init file * fixed typeguard error and usage of private/experimental api. * restructured wrapper serialization and removed unnecessary components. * expanded on docstr and added repr * cleaned up docstrings, added assertion tests, and added explicit test for only the serialization * ran black and flake8 * fixed doc string Co-authored-by: gabrieldemarmiesse <gabrieldemarmiesse@gmail.com>

hyang0129 added 6 commits January 28, 2020 21:06

initial setup. need to build tests

e2c1378

build some tests. need to test them

ef4b235

fixed typo

34e4e16

created first test

c9e8b99

created first test

3b545a4

accidentally messed up another file

fab5871

hyang0129 requested review from facaiy, Squadrick and WindQAQ as code owners January 29, 2020 18:10

boring-cyborg bot added the optimizers label Jan 29, 2020

This comment has been minimized.

Sign in to view

googlebot added the cla: no label Jan 29, 2020

googlebot added cla: yes and removed cla: no labels Jan 29, 2020

hyang0129 added 13 commits January 29, 2020 14:02

accidentally messed up another file

333e907

accidentally messed up another file

b79135d

added run all distributed

c4d2588

fixed formatting

c47c34a

trying to fix tests not running on github CI.

a510f70

realized that I should probably add the new optimizer files to the bu…

6940433

…ild and init

added typeguard and docstring

af932d8

removed run_all_distributed

ea1f62e

graph and eager testing for SGD

80e2768

reformatted

f89ac90

added distributed tests

1137bc3

removed distributed tests

1149719

reverted discriminative layer grad adjust back to apply gradients

87b396c

WindQAQ reviewed Sep 10, 2020

View reviewed changes

restructured wrapper serialization and removed unnecessary components.

35dca9b

expanded on docstr and added repr

f189c39

WindQAQ reviewed Sep 14, 2020

View reviewed changes

hyang0129 added 3 commits September 13, 2020 21:26

cleaned up docstrings, added assertion tests, and added explicit test…

c63c635

… for only the serialization

ran black and flake8

2c46752

fixed doc string

4b856f1

hyang0129 requested a review from WindQAQ September 14, 2020 17:35

WindQAQ approved these changes Sep 14, 2020

View reviewed changes

WindQAQ merged commit 3897e2f into tensorflow:master Sep 14, 2020

This was referenced Sep 14, 2020

Fixing minor issues from #969 #2163

Closed

fixed typo in multioptimizer class name and added code owners #2164

Merged

WindQAQ mentioned this pull request Sep 25, 2020

Fix MultiOptimizer list of layers #2180

Merged

21 tasks



		@tf.keras.utils.register_keras_serializable(package="Addons")
		class MultiOpt(Optimizer):

		@@ -0,0 +1,101 @@
		# Copyright 2019 The TensorFlow Authors. All Rights Reserved.



		@tf.keras.utils.register_keras_serializable(package="Addons")
		class MultiOptimzer(Optimizer):


		@classmethod
		def create_optimizer_spec(cls, optimizer_instance, layer):

Discriminative Layer Training #969

Discriminative Layer Training #969

Conversation

hyang0129 commented Jan 29, 2020

This comment has been minimized.

hyang0129 commented Jan 29, 2020

hyang0129 commented Jan 29, 2020

googlebot commented Jan 29, 2020

hyang0129 commented Sep 9, 2020

WindQAQ left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WindQAQ Sep 13, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

WindQAQ commented Sep 10, 2020

hyang0129 commented Sep 13, 2020

hyang0129 commented Sep 13, 2020

WindQAQ commented Sep 13, 2020

hyang0129 commented Sep 14, 2020

WindQAQ left a comment

Choose a reason for hiding this comment

WindQAQ Sep 14, 2020 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hyang0129 commented Sep 14, 2020

WindQAQ left a comment

Choose a reason for hiding this comment

hyang0129 commented Sep 14, 2020

WindQAQ commented Sep 14, 2020

WindQAQ commented Sep 14, 2020

hyang0129 commented Sep 14, 2020

WindQAQ commented Sep 14, 2020

hyang0129 commented Sep 14, 2020

seanpmorgan commented Sep 16, 2020

WindQAQ Sep 13, 2020 •

edited

Loading

WindQAQ Sep 14, 2020 •

edited

Loading