Model.fit 使用 Colab 进行训练需要很长时间

问：

我正在尝试创建一个神经网络来求解泊松偏微分方程 （PDE）。但是，由于某种原因，处理纪元需要很长时间，平均需要 7 分钟。我不知道问题可能出在哪里。有人可以帮我吗？

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers

model = keras.Sequential([
    keras.layers.Input(shape=(2,)),
    keras.layers.Dense(10, activation='sigmoid'),
    keras.layers.Dense(1, activation='linear')
])

def A(x, y):
    f0_y = tf.zeros_like(y)
    f1_y = tf.zeros_like(y)
    g0_x = tf.sin(np.pi * x)
    g1_x = tf.sin(np.pi * x)

    return (1 - x) * f0_y + x * f1_y + (1 - y) * (g0_x - ((1 - x) * g0_x + x * g0_x)) + y * (g1_x - ((1 - x) * g1_x + x * g1_x))

def u_t(x, y, model):
    N = model(tf.stack([x, y], axis=1))
    return A(x, y) + x * (1 - x) * y * (1 - y) * N

num_points = 100  # Número de pontos no domínio
x = np.linspace(0, 1, num_points)
y = np.linspace(0, 1, num_points)
X, Y = np.meshgrid(x, y)
X_flat = X.flatten()
Y_flat = Y.flatten()

def custom_loss(model, x, y):
    x = tf.convert_to_tensor(x, dtype=tf.float32)
    y = tf.convert_to_tensor(y, dtype=tf.float32)

    with tf.GradientTape() as tape2:
        tape2.watch(x)
        tape2.watch(y)
        with tf.GradientTape() as tape:
            tape.watch(x)
            tape.watch(y)
            u = u_t(x, y, model)
        du_dx, du_dy = tape.gradient(u, [x, y])
    d2u_dx2, d2u_dy2 = tape2.gradient([du_dx, du_dy], [x, y])

    laplacian_u = d2u_dx2 + d2u_dy2

    f = -5 * np.pi**2 * tf.sin(np.pi * x) * tf.cos(2 * np.pi * y)

    return tf.reduce_mean(tf.square(laplacian_u - f))

input_data = np.stack([X_flat, Y_flat], axis=1)

model.compile(optimizer='adam', loss=lambda y_true, y_pred: custom_loss(model, X_flat, Y_flat))

model.fit(x=input_data, y=np.zeros_like(X_flat), epochs=100, batch_size=32)