dfdx_diff_net = lambda x: diff_net(x)[1]
d2fd2x_diff_net = lambda x: diff_net(x, hessian=True)[2]
with torch.no_grad():
t0_jac = time.perf_counter()
dydx_hat = dfdx_diff_net(x_torch)
t_for_jac = time.perf_counter() - t0_jac
dydx_hat = dydx_hat.detach().cpu().numpy().squeeze()
t0_jac = time.perf_counter()
d2yd2x_hat = d2fd2x_diff_net(x_torch)
t_for_hes = time.perf_counter() - t0_jac
d2yd2x_hat = d2yd2x_hat.detach().cpu().numpy().squeeze()
t0_jac = time.perf_counter()
dydx_autograd[i] = jacobian(f_diff_net, x_torch, create_graph=False)
t_rev_jac = time.perf_counter() - t0_jac
dydx_autograd[i] = dydx_autograd[i].detach().cpu().numpy().squeeze()
t0_hes = time.perf_counter()
d2yd2x_autograd[i] = jacobian(dfdx_diff_net, x_torch, create_graph=False, v1=False)
t_rev_hes = time.perf_counter() - t0_hes
d2yd2x_autograd[i] = d2yd2x_autograd[i].detach().cpu().numpy().squeeze()
seq_net = nn.Sequential(nn.Linear(1, hyper['n_width']), nn.Tanh(), nn.Linear(hyper['n_width'], 1)).cuda()
t0_auto = time.perf_counter()
dydx_auto_test = jacobian_auto(seq_net, x_torch, create_graph=True)
t_test = time.perf_counter() - t0_auto
dydx_auto_test = dydx_auto_test.detach().cpu().numpy().squeeze()
err_dydx = 1. / float(x_test.shape[0]) * np.sum((dydx_test - plot_dydx_test[i]) ** 2)
with torch.no_grad():
t0_jac = time.perf_counter()
dydx_hat = dfdx_diff_net(x_torch)
t_for_jac = time.perf_counter() - t0_jac
dydx_hat = dydx_hat.detach().cpu().numpy().squeeze()
t0_jac = time.perf_counter()
d2yd2x_hat = d2fd2x_diff_net(x_torch)
t_for_hes = time.perf_counter() - t0_jac
d2yd2x_hat = d2yd2x_hat.detach().cpu().numpy().squeeze()
t0_jac = time.perf_counter()
dydx_autograd[i] = jacobian(f_diff_net,
x_torch,
create_graph=False,
v1=False).transpose(0, 1).view(
hyper['n_ensemble'], -1)
t_rev_jac = time.perf_counter() - t0_jac
dydx_autograd[i] = dydx_autograd[i].detach().cpu().numpy().squeeze()
t0_hes = time.perf_counter()
d2yd2x_autograd[i] = jacobian(dfdx_diff_net,
x_torch,
create_graph=False,
v1=False).transpose(0, 1).view(
hyper['n_ensemble'], -1)
t_rev_hes = time.perf_counter() - t0_hes
d2yd2x_autograd[i] = d2yd2x_autograd[i].detach().cpu().numpy().squeeze(
)
plot_y_test[i] = y_hat
plot_dydx_test[i] = dydx_hat
print("\n################################################")
print("Autograd Performance:")
## Autograd Test:
f_diff_net = lambda x: diff_net(x)[0]
dfdx_diff_net = lambda x: diff_net(x)[1]
t0_jac = time.perf_counter()
dydx_hat = dfdx_diff_net(x_torch).detach().cpu().numpy().squeeze()
t_for_jac = time.perf_counter() - t0_jac
t0_jac = time.perf_counter()
dydx_autograd[i] = jacobian(f_diff_net, x_torch, create_graph=False).detach().cpu().numpy().squeeze()
t_rev_jac = time.perf_counter() - t0_jac
t0_hes = time.perf_counter()
d2yd2x_autograd[i] = jacobian(dfdx_diff_net, x_torch, create_graph=False, v1=False).detach().cpu().numpy().squeeze()
t_rev_hes = time.perf_counter() - t0_hes
err_dydx = 1. / float(x_test.shape[0]) * np.sum((dydx_test - plot_dydx_test[i]) ** 2)
err_d2yd2x = 1. / float(x_test.shape[0]) * np.sum((d2yd2x_test.squeeze() - d2yd2x_autograd[i]) ** 2)
err_autograd = 1. / float(x_test.shape[0]) * np.sum((dydx_hat - dydx_autograd[i]) ** 2)
print(f"Jacobian:\n"
f" Approximation MSE = {err_dydx:.3e}\n"
f" Diff Mode MSE = {err_autograd:.3e}\n"
f" Reverse Diff = {t_rev_jac:.3e}s\n"
f" Forward Diff = {t_for_jac:.3e}s\n"
return out
def jacobian_fun(input):
input = input.unsqueeze(0) if input.ndim == 2 else input
out = torch.matmul(A + A.transpose(dim0=1, dim1=2), input)
return out
def hessian_fun(input):
out = (A + A.transpose(dim0=1, dim1=2)).repeat(input.shape[0], 1, 1)
return out
y, dydx, d2yd2x = fun(x), jacobian_fun(x), hessian_fun(x)
t0 = time.perf_counter()
for i in range(n_iter):
auto_dydx = jacobian(fun, x)
t_jacobian = (time.perf_counter() - t0) / n_iter
t0 = time.perf_counter()
for i in range(n_iter):
auto_d2yd2x = jacobian(jacobian_fun, dydx)
t_hessian = (time.perf_counter() - t0) / n_iter
t0 = time.perf_counter()
for i in range(n_iter):
auto_d2yd2x = hessian(fun, dydx)
t_auto_hessian = (time.perf_counter() - t0) / n_iter
t0 = time.perf_counter()
for i in range(n_iter):