-
Notifications
You must be signed in to change notification settings - Fork 25
/
custom_grad_example3.py
97 lines (69 loc) · 2.6 KB
/
custom_grad_example3.py
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
#*************************************************************************
# Copyright 2021 Adobe Systems Incorporated.
#
# Please see the attached LICENSE file for more information.
#
#**************************************************************************/
# Custom gradient example with vector (signal) and scalar arguments of the function:
# f(x, y) = x^2 * y, where x is a vector
import numpy as np
import tensorflow as tf
# Python function with no custom gradient
def foo_no_grad(x, y):
y = np.square(x)*y[0]#*y[1]
return tf.constant(y)
# Python function with custom analytic gradient
@tf.custom_gradient
def foo_custom_grad_analytic(x, y):
z = np.square(x)*y#*y[1]
def grad_fn(dy):
# Gradient w.r.t x
gradx = 2 * np.array(x) * y
# Gradient w.r.t y
grady = np.square(x)
# Returns vector-Jacobian
return gradx * dy, np.dot(np.transpose(dy), grady)
return z, grad_fn
# Python function with custom numerical gradient
@tf.custom_gradient
def foo_custom_grad_numeric(x, y):
epsilon = 0.01
def func(x, y):
z = np.square(x)*y
return z
def grad_fn(dy):
# Grad w.r.t x
J_plus = func(x + epsilon, y)
J_minus = func(x - epsilon, y)
gradx = (J_plus - J_minus)/(2.0*epsilon)
# Grad w.r.t y
J_plus = func(x, epsilon + y)
J_minus = func(x, - epsilon + y)
grady = (J_plus - J_minus)/(2.0*epsilon)
return gradx * dy, np.dot(np.transpose(dy), grady)
return func(x, y), grad_fn
# Python function only using tf ops with gradient automatically defined via auto-diff
def foo_autodiff(x, y):
y = y*tf.square(x)
return y
with tf.GradientTape(persistent=True) as tape:
a = 2*np.ones((3,1)) # signal
b = 3*np.ones((1,1)); #b[1] = 1 # params
x = tf.constant(a, dtype=tf.float64)
y = tf.constant(b, dtype=tf.float64)
tape.watch(x)
tape.watch(y)
z3 = foo_autodiff(x, y)**2
z1 = foo_no_grad(x,y)**2
z2 = foo_custom_grad_analytic(x, y)**2
z4 = foo_custom_grad_numeric(x, y)**2
print('\nGrad w.r.t. x')
print("foo_no_grad", tape.gradient(z1, x))
print("foo_custom_grad_analytic", tape.gradient(z2, x))
print("foo_autodiff", tape.gradient(z3, x))
print("foo_custom_grad_numeric", tape.gradient(z4, x))
print('\nGrad w.r.t. y')
print("foo_no_grad", tape.gradient(z1, y))
print("foo_custom_grad_analytic", tape.gradient(z2, y))
print("foo_autodiff", tape.gradient(z3, y))
print("foo_custom_grad_numeric", tape.gradient(z4, y))