DeepESM/flipGradientTF2.py at master · StealthStop/DeepESM · GitHub

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
import tensorflow as tf

#@tf.custom_gradient
#def GradientReversalOperator(x, hp_lambda):
#        def grad(dy):
#                return -hp_lambda * dy
#        return x, grad
#
#class GradientReversal(tf.keras.layers.Layer):
#        def __init__(self, hp_lambda=1.0):
#                super(GradientReversal, self).__init__()
#                self.hp_lambda = hp_lambda
#
#        def call(self, x):
#                return GradientReversalOperator(x, self.hp_lambda)


@tf.custom_gradient
def GradientReversalOperator(x):
        def grad(dy):
                return -1 * dy
        return x, grad

class GradientReversal(tf.keras.layers.Layer):
        def __init__(self):
                super(GradientReversal, self).__init__()
                self.hp_lambda = 1.0

        def setLambda(self, hp_lambda):
                self.hp_lambda = hp_lambda

        def call(self, inputs):
                return GradientReversalOperator(inputs)