Hi, I am trying to implement some code from Tensorflow, but I am uncertain on how to implement their method of backpropogation. Here is their code:
self.IRK_weights = np.reshape(tmp[0:q**2+q], (q+1,q))
self.x0_tf = tf.placeholder(tf.float32, shape=(None, self.x0.shape[1]))
self.u0_tf = tf.placeholder(tf.float32, shape=(None, self.u0.shape[1]))
self.dummy_x0_tf = tf.placeholder(tf.float32, shape=(None, self.q)) #dummy variable for fwd_gradients
def fwd_gradients_0(self, U, x):
g = tf.gradients(U, x, grad_ys=self.dummy_x0_tf)[0]
return tf.gradients(g, self.dummy_x0_tf)[0]
def net_U0(self, x):
nu = 0.01/np.pi
U1 = self.neural_net(x, self.weights, self.biases)
U = U1[:,:-1]
U_x = self.fwd_gradients_0(U, x)
U_xx = self.fwd_gradients_0(U_x, x)
F = -U*U_x + nu*U_xx
U0 = U1 - self.dt*tf.matmul(F, self.IRK_weights.T)
return U0
The key part that I am struggling to understand is where they calculate the x derivative of u and use the function fwd_gradients_0, where they also use a dummy variable. Here is how I tried to implement it which is seemingly wrong:
x.requires_grad_(True)
s = nn.Tanh()(self.fc1(x.float()))
s = nn.Tanh()(self.fc2(s))
s = nn.Tanh()(self.fc3(s))
u_1 = self.fc4(s)
u = u_1[:, :-1]
u_x = autograd.grad(u, x, torch.ones(u.shape).to(device), create_graph=True)[0]
u_xx = autograd.grad(u_x, x, torch.ones(u_x.shape).to(device), create_graph=True)[0]
f = - u * u_x + (0.01/np.pi)*u_xx
u_0 = u_1 - dt * torch.mm(f.double(), IRK_weights.T.double())
Any help would be greatly appreciated to help me understand/solve this!