Custom Neural Network Architectures¶

Modification of Neural Network Architecture¶

By default, TensorDiffEq will build a fully-connected network using the layer sizes and lengths you define in the layer_sizes parameter, which is fed into the model.compile call. However, once the model has been compiled, that network can be overwritten with any Keras neural network. Here we demonstrate how to do so, adding batch norm layers to the network. Referencing the example here, we can modify the neural network as such:

# need to include keras.layers and Sequential API
from tf.keras import layers, Sequential

layer_sizes = [2, 128, 128, 128, 128, 1]

model_bn = tf.keras.Sequential(
    [
        layers.Dense(2, activation=tf.nn.tanh,
            kernel_initializer="glorot_normal"),        
            layers.BatchNormalization(),
            layers.Dense(128, activation=tf.nn.tanh,
            kernel_initializer="glorot_normal"), 
            layers.BatchNormalization(),
            layers.Dense(128, activation=tf.nn.tanh,
            kernel_initializer="glorot_normal"),
            layers.BatchNormalization(),           
            layers.Dense(1, activation=None),
    ]
)

model = CollocationSolverND()
model.compile(layer_sizes, f_model, Domain, BCs)
# overwrite the default NN with our new one defined above
model.u_model = model_bn
model.fit(tf_iter=1000, newton_iter=1000)

This will fit your custom network (i.e., with batch norm) as the PDE approximation network, allowing more stability and reducing the likelihood of vanishing gradients in the training.