Output Transform with GPU

[LordPontiag]

Hi
When I want to apply hard constraint in this way:

@tf.function
def output_transform(x, y):
    a      = y[:, 0:1]
    b      = y[:, 1:2]
    c      = y[:, 2:3]
    d      = tf.maximum(y[:, 3:4], 1e-4)
    e      = tf.maximum(y[:, 4:5], 1e-4)
    return tf.concat([a,b,c,d,e], axis=1)

net = dde.nn.FNN(layer_size, activation, initializer)
net.apply_output_transform(output_transform)

it doesn't switch on GPU.
Without this output_transform, it runs on GPU.
I even tried tf.nn.softplus, there was same issue with some errors.
What is the problem?!

[vl-dud]

Have you tried to remove the decorator @tf.function?

[LordPontiag]

Entire code is about 2000 lines!. But I can provide some:
Libraries:

import os
import time
import Ofpp as Myfoam
import deepxde as dde
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.tri as tri
from matplotlib import cm
import pyvista as pv
import pandas as pd
import tensorflow as tf
import scipy.interpolate as interp
from scipy.spatial import cKDTree
from sklearn.model_selection import train_test_split

dde.backend.set_default_backend("tensorflow")

# Verify the policy is set
print("Mixed precision policy:", tf.keras.mixed_precision.global_policy())

# Function to check if GPU is available and set the device accordingly
def use_gpu(gpu=True):
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            if gpu:
                # Set memory growth for each GPU to avoid OOM errors
                for gpu in gpus:
                    tf.config.experimental.set_memory_growth(gpu, True)
                print("Using GPU")
                # Print GPU device name
                for i, gpu in enumerate(gpus):
                    print(f"GPU {i}: {gpu.name}")
            else:
                print("Switching to CPU")
        except RuntimeError as e:
            # Memory growth must be set before GPUs have been initialized
            print(e)
    else:
        print("No GPU devices found. Using CPU.")

# tf.compat.v1.enable_eager_execution()
use_gpu()

print("Num GPUs Available: ", len(tf.config.experimental.list_physical_devices('GPU')))
print("TensorFlow version:", tf.__version__)

Geometry and BC:

if cond:
    geom = dde.geometry.PointCloud(
                                    points=all_coords_3D,
                                    boundary_points=boundary_coords_3D,
                                    boundary_normals=boundary_normals_3D
                                    )
    BC_U     = dde.icbc.PointSetBC(wall_coords_3D[::10]  , 0, component=0) 
    BC_V     = dde.icbc.PointSetBC(wall_coords_3D[::10]  , 0, component=1)  
    BC_P     = dde.icbc.PointSetBC(wall_coords_3D[::10]  , 0, component=3) 

PDEs:

def pde(X, Y):

   # Navier_Stokes

    # Extract variables
    a = Y[:, 0:1]  # Velocity in x-direction
    b = Y[:, 1:2]  # Velocity in y-direction
    c = Y[:, 2:3]  # Pressure

    # Gradients for U, V, P, k, omega
    a_x = dde.grad.jacobian(Y, X, i=0, j=0)
    a_y = dde.grad.jacobian(Y, X, i=0, j=1)
    b_x = dde.grad.jacobian(Y, X, i=1, j=0)
    b_y = dde.grad.jacobian(Y, X, i=1, j=1)
   .
   .
   .
    f1 = U_x + V_y
   .
   .
   .
    return [f1, f2, f3, f4, f5]

NN Setup and train:

@tf.function
def output_transform(x, y):
    a      = y[:, 0:1]
    b      = y[:, 1:2]
    c      = y[:, 2:3]
    d      = tf.maximum(y[:, 3:4], 1e-4)
    e      = tf.maximum(y[:, 4:5], 1e-4)
    return tf.concat([a,b,c,d,e], axis=1)

# Single Run
layers   = 7  # Number of layers
neurons  = 49  # Neurons per layer

# Define the network architecture
layer_size  = [NN_input] + [neurons] * layers + [5]
activation  = "tanh"
initializer = "Glorot uniform"

# Re-initialize the network for each test case
net = dde.nn.FNN(layer_size, activation, initializer)        
net.apply_output_transform(output_transform)                # The Problematic Line, which prevents switching to GPU

# Main training
data = dde.data.PDE(
    geom,
    pde,
    [BC],
    train_distribution= "uniform",
    anchors= Training_POINTS,
    num_test= 7500)

weights = [1,1,1,1,1,
                    1,1,1,1,
                    1,1,1,1,1]

iterations = 10000
model = dde.Model(data, net)
model.compile("adam", lr=1e-4, loss_weights=weights, external_trainable_variables=exteranl_VARs)

# Train the model with Adam
losshistory, train_state = model.train(
                iterations = iterations, 
                display_every = 1000)

# Train the model with L-BFGS
dde.optimizers.config.set_LBFGS_options(maxiter=10000)
model.compile("L-BFGS")
losshistory, train_state = model.train(display_every=10000)

[vl-dud]

Can you show a minimal version of the code that reproduces the problem? So that I can run the code and see it.

[LordPontiag]

Thank you for your time and help. But it takes a long time to do this. I'm not sure, but I reckon the problem is with the backend of TensorFlow 2X. Because everything works great until the net.apply_output_transform(output_transform) becomes activated. It works after that, but not on the GPU. It may have limitations in its development. @lululxvi , Professor Lu, do you have any idea in this regard?

[LordPontiag]

@vl-dud Hi,
I finally found the problem's solution!
The problem of "not using GPU" while "using output transform" is because of the TensorFlow 2X !!!
If you use TensorFlow 1X (i.e. tf.compat.v1), there wouldn't be such problem!