如何将含自定义损失函数的TensorFlow 1.x代码迁移至TensorFlow 2.x

转换后完整代码 #

import DGM
import tensorflow as tf
import numpy as np
import scipy.stats as spstats
import matplotlib.pyplot as plt
from tqdm.notebook import trange

# Option parameters
phi = 10
n = 0.01
T = 4

# Solution parameters (domain on which to solve PDE)
t_low = 0.0 - 1e-10
x_low = 0.0 + 1e-10 
x_high = 1.0

# neural network parameters
num_layers = 3
nodes_per_layer = 50

# Training parameters
sampling_stages  = 2500  # number of times to resample new time-space domain points
steps_per_sample = 20    # number of SGD steps to take before re-sampling

# Sampling parameters
nsim_interior   = 100
nsim_boundary_1 = 50
nsim_boundary_2 = 50
nsim_initial    = 50
x_multiplier    = 1.1    # multiplier for oversampling i.e. draw x from [x_low, x_high * x_multiplier]


def sampler(nsim_interior, nsim_boundary_1, nsim_boundary_2, nsim_initial):
    ''' Sample time-space points from the function's domain; points are sampled
        uniformly on the interior of the domain, at the initial/terminal time points
        and along the spatial boundary at different time points. 
    
    Args:
        nsim_interior:      number of space points in the interior of U
        nsim_boundary_1:    number of space points in the boundary of U
        nsim_boundary_2:    number of space points in the boundary of U_x
        nsim_initial:       number of space points at the initial time
    ''' 
    
    # Sampler #1: domain interior
    t_interior = np.random.uniform(low=t_low, high=T, size=[nsim_interior, 1])
    x_interior = np.random.uniform(low=x_low, high=x_high*x_multiplier, size=[nsim_interior, 1])

    # Sampler #2: spatial boundary 1
    t_boundary_1 = np.random.uniform(low=t_low, high=T, size=[nsim_boundary_1, 1])
    x_boundary_1 = np.ones((nsim_boundary_1, 1))

    # Sampler #3: spatial boundary 2
    t_boundary_2 = np.random.uniform(low=t_low, high=T, size=[nsim_boundary_2, 1])
    x_boundary_2 = np.zeros((nsim_boundary_2, 1))
    
    # Sampler #4: initial condition
    t_initial = np.zeros((nsim_initial, 1))
    x_initial = np.random.uniform(low=x_low, high=x_high*x_multiplier, size=[nsim_initial, 1])
    
    return (
        t_interior, x_interior,
        t_boundary_1, x_boundary_1,
        t_boundary_2, x_boundary_2,
        t_initial, x_initial
    )


def compute_loss(
    model,
    t_interior, x_interior,
    t_boundary_1, x_boundary_1,
    t_boundary_2, x_boundary_2,
    t_initial, x_initial
    ):
    ''' Compute total loss for training.
    
    Args:
        model:                          DGM model object
        t_interior, x_interior:         sampled time / space points in the interior of U
        t_boundary_1, x_boundary_1:     sampled time / space points in the boundary of U
        t_boundary_2, x_boundary_2:     sampled time / space points in the boundary of U_x
        t_initial, x_initial:           sampled time / space points at the initial time
    '''  

    # Loss term #1: PDE
    with tf.GradientTape(persistent=True) as tape:
        tape.watch([t_interior, x_interior, x_boundary_2])
        # 计算内部点的函数值与一阶导数
        u = model(t_interior, x_interior)
        u_t = tape.gradient(u, t_interior)
        u_x = tape.gradient(u, x_interior)
        # 计算二阶导数
        u_xx = tape.gradient(u_x, x_interior)
        
        # 计算边界2的一阶导数
        u_boundary2 = model(t_boundary_2, x_boundary_2)
        u_x_boundary2 = tape.gradient(u_boundary2, x_boundary_2)
    
    # 删除persistent tape释放资源
    del tape
    
    diff_u = u_t - u_xx + phi**2 * (tf.nn.relu(u) + 1e-10)**n
    L1 = tf.reduce_mean(tf.square(diff_u))
    
    # Loss term #2: First boundary condition
    u_boundary1 = model(t_boundary_1, x_boundary_1)
    bc1_error = u_boundary1 - 1

    # Loss term #3: Second boundary condition
    bc2_error = u_x_boundary2 - 0

    # Loss term #4: Initial condition
    u_initial = model(t_initial, x_initial)
    init_error = u_initial - 1

    # 计算边界与初始条件的平均L2损失
    L2 = tf.reduce_mean(tf.square(bc1_error + bc2_error + init_error))

    return L1 + L2, L1, L2


# initialize DGM model (last input: space dimension = 1)
model = DGM.DGMNet(nodes_per_layer, num_layers, 1)

# 设置学习率调度与优化器
starting_learning_rate = 3e-4
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=starting_learning_rate,
    decay_steps=100000,
    decay_rate=0.96,
    staircase=True
)
optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)

# 加载检查点
try:
    model.load_weights("checkpoint/")
    print("Loading from checkpoint.")
except:
    print("Checkpoint not found.")

# 训练循环
for i in trange(sampling_stages):
    # 采样数据
    t_interior, x_interior, \
    t_boundary_1, x_boundary_1, \
    t_boundary_2, x_boundary_2, \
    t_initial, x_initial = sampler(
        nsim_interior, nsim_boundary_1, nsim_boundary_2, nsim_initial
    )
    
    # 每个样本执行多步SGD
    for _ in range(steps_per_sample):
        with tf.GradientTape() as tape:
            total_loss, L1, L2 = compute_loss(
                model,
                t_interior, x_interior,
                t_boundary_1, x_boundary_1,
                t_boundary_2, x_boundary_2,
                t_initial, x_initial
            )
        
        # 计算梯度并更新参数
        grads = tape.gradient(total_loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
    
    # 定期打印损失
    if i % 10 == 0:
        print(f"Loss: {total_loss.numpy():.5f},\t L1: {L1.numpy():.5f},\t L2: {L2.numpy():.5f},\t iteration: {i}")

    # 保存检查点
    model.save_weights("checkpoint/")

关键转换说明 #

• 移除旧API依赖：完全删除tf.compat.v1.placeholder、tf.compat.v1.Session等旧版接口，直接使用numpy数组作为模型输入，依托TF2.x即时执行模式运行。
• 自动微分重构：用tf.GradientTape替代tf.gradients，通过persistent=True支持多次梯度计算（用于求解二阶偏导），计算完成后手动删除tape释放资源。
• 损失函数适配：保留原损失的物理逻辑，将梯度计算逻辑迁移到GradientTape上下文内，确保自动微分能正确追踪模型参数的变化。
• 优化器与调度升级：替换为TF2.x原生的tf.keras.optimizers.Adam和ExponentialDecay学习率调度，对齐Keras生态的使用方式。
• 训练循环重构：删除会话运行逻辑，改为在梯度带内计算损失、获取梯度，再通过优化器手动更新模型参数，符合TF2.x的命令式编程风格。

内容的提问来源于stack exchange，提问作者Arman Bolatov