双神经网络单损失函数如何自定义优化

Question

毕设课题，涉及到PINNs（物理信息神经网络），简单来说就是损失函数可以直接通过模型输出与输入数据进行计算来得到，所以训练时不需要提供准确答案，直接随机取点就好了。（用来解复杂偏微分方程的，因为当偏微分方程很复杂时，直接解方程难度极大）。
现在需要的模型是自定义损失函数和训练过程，但是需要用到两个不同的DNN，简单来说就是两个输入，x和y，x进一个DNN，y进另一个DNN，出来的两个结果组成损失函数（比如loss=(x_pred-x^2)+(y_pred-y^3)）。
机械类专业，中途出家当程序员，神经网络刚学
我使用的单DNN模型如下，可以的话，希望能就我的自定义训练框架进行更改，方便我理解（因为简单）。

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

X = tf.constant([[1.0], [2.0], [3.0], [4.0]])

model = keras.Sequential([
    keras.layers.Dense(1),
])
# 优化器
optimizer = tf.keras.optimizers.Adam(learning_rate=0.1, epsilon=1e-08)


for i in range(200):
    with tf.GradientTape(persistent=True) as tape:
        tape.watch(X)
        y_pred = model(X)
        Y = X * X
        y_exac = tape.gradient(Y, X)
        loss = tf.reduce_mean(tf.square(y_pred - y_exac))
    grads = tape.gradient(loss, model.variables)
    # 更新参数
    optimizer.apply_gradients(grads_and_vars=zip(grads, model.variables))

print(model.predict([5.0]))

如果的话，可以举output=x^2+y^3的例子

Answer 1

一定要用tensorflow吗