Now lets take some of these examples, make them a little “larger” (by adding more collocation points, for instance), and solve the across multiple GPUs.
An extremely unique feature of TensorDiffEq is that the exact same code that implements a solver on a small scale CPU platform can implement a solver on a massive scale. The only major difference is scaling up the number of collocation points and implementing a distributed solver in TensorDiffEq.
Even more powerful - this only requires the modification of one line of code - a single boolean value to be modified. See the below Allen-Cahn example:
Domain = DomainND(["x", "t"], time_var='t') Domain.add("x", [-1.0, 1.0], 512) Domain.add("t", [0.0, 1.0], 201) N_f = 1000000 # 1m collocation points Domain.generate_collocation_points(N_f) def func_ic(x): return x ** 2 * np.cos(math.pi * x) # Conditions to be considered at the boundaries for the periodic BC def deriv_model(u_model, x, t): u = u_model(tf.concat([x, t], 1)) u_x = tf.gradients(u, x) u_xx = tf.gradients(u_x, x) return u, u_x init = IC(Domain, [func_ic], var=[['x']]) x_periodic = periodicBC(Domain, ['x'], [deriv_model]) BCs = [init, x_periodic] def f_model(u_model, x, t): u = u_model(tf.concat([x, t], 1)) u_x = tf.gradients(u, x) u_xx = tf.gradients(u_x, x) u_t = tf.gradients(u, t) c1 = tdq.utils.constant(.0001) c2 = tdq.utils.constant(5.0) f_u = u_t - c1 * u_xx + c2 * u * u * u - c2 * u return f_u layer_sizes = [2, 128, 128, 128, 128, 1] model = CollocationSolverND() model.compile(layer_sizes, f_model, Domain, BCs, dist=True) model.fit(tf_iter=1000)
We note here that the only difference in the code comes in the
compile call, where we add the argument
This enables TensorDiffEq to adopt a
data-parallelism approach to training. In this case, the collocation points are distributed evenly across all available workers. This is one of
the most powerful aspects of TensorDiffEq, being able to scale readily without modification of the code. The same physics model, neural network model, and
optimizer (except for L-BFGS, at the time of this writing) can be run on a small model on a local machine, and can also be scaled up to run on an
enterprise-level data center with \(N\) GPUs.
Notes and Best-Practices for GPUs¶
graph-mode L-BFGS is typically faster on a single-GPU, given sufficient model size. On a CPU, empirically, it has been demonstrated that eager-mode L-BFGS is actually faster than graph-mode
We are interested in community feedback! If you notice something interesting here, open a PR on these docs and let us know!