Multi-GPU Support Assistance

[helman44]

Hello All,

I have been reading through the discussions and found that multi-GPU support was enabled via tf.distribute on Sionna v0.16.0. Self admittedly I'm a PyTorch user and haven't delved into TensorFlow enough to easily implement this strategy in my project. I'm currently simulating an aerial path over a target and using one GPU out of the two A6000s I have available. I'm also checking the compute time for a number of max bounce depths. Please see the current implementation below (apologies if this isn't formatted correctly):

import os
gpu_num = 0 # Use "" to use the CPU
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import mitsuba as mi
try:
    import google.colab
    colab_compat = True # deactivate preview
except:
       colab_compat = False
resolution = [480,320] # increase for higher quality of renderings

class ExitCell(Exception):
    def _render_traceback_(self):
        pass

try:    import sionna
except ImportError as e:
 
   import os
   os.system("pip install sionna")
   import sionna

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_memory_growth(gpus[0], True)
    except RuntimeError as e:
        print(e)

tf.get_logger().setLevel('ERROR')

tf.random.set_seed(1) # Set global random seed for reproducibility

mirrored_strategy = tf.distribute.MirroredStrategy()

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import sys
from tqdm import tqdm

from sionna.rt import load_scene, PlanarArray, Transmitter, Receiver, RadioMaterial, Camera
from sionna.rt.utils import r_hat
from sionna.constants import PI, SPEED_OF_LIGHT
from sionna.utils import expand_to_rank
from scipy.io import savemat
from numpy import *
import math
import time

print(gpus)
scene = load_scene("slicy.xml")

aperture_start = 0
aperture_stop = 1
CPI_length = np.abs(aperture_stop - aperture_start) * 256
elevation = 15
ant_range = 50

obj = scene.get("slicy")
obj.radio_material = "itu_concrete" # "slicy" is made of "concrete"


scene.tx_array = PlanarArray(num_rows=1,
                             num_cols=1,
                             vertical_spacing=0.5,
                             horizontal_spacing=0.5,
                             pattern="iso",
                             polarization="V")

scene.rx_array = scene.tx_array

scene.add(Transmitter(name="tx", position=[0,0,0]))
scene.add(Receiver(name="rx", position=[0,0,0]))

x_pos = 0 - (ant_range * sin(radians(elevation))) * cos(radians(np.linspace(aperture_start, aperture_stop, CPI_length)))
y_pos = 0 - (ant_range * sin(radians(elevation))) * sin(radians(np.linspace(aperture_start, aperture_stop, CPI_length)))
z_pos = np.ones(CPI_length)*(ant_range * sin(radians(elevation)))

radar_trajectory = np.array([x_pos, y_pos, z_pos]).T 

max_depths = 10 # evaluate performance up to 10 reflections
depths = range(1,max_depths+1)
ts = []
pathlist = [dict() for x in range(CPI_length)]
for d in depths:
    with tqdm(range(CPI_length)) as pbar:
        # save start time
        t = time.time()
        for pulse in pbar:
            scene.transmitters["tx"].position = radar_trajectory[pulse]
            scene.receivers["rx"].position = radar_trajectory[pulse] 
            scene.transmitters["tx"].look_at([0,0,0])
            scene.receivers["rx"].look_at([0,0,0]) 
            paths = scene.compute_paths(los=False, 
                                        reflection=True, 
                                        scattering=True,
                                        scat_random_phases=True,
                                        scat_keep_prob=0.001,
                                        num_samples=1e6,
                                        max_depth=d)
             paths.normalize_delays=False
             pathlist[pulse] = paths.to_dict()
    ts.append(time.time()-t)

[merlinND]

Hello @helman44,

Unfortunately, the ray tracing part of Sionna does not support multi-GPU due to missing support in Mitsuba (the underlying library used for ray tracing, which itself uses OptiX).
The easiest way would be to start your script twice and split the work over e.g. depths, CPI_LENGTH, or any other dimension that makes sense for your problem.
Something to the effect of:

CUDA_VISIBLE_DEVICES=0 python ./run.py --depths 0 1 2 3 4 5 6

# In another terminal
CUDA_VISIBLE_DEVICES=1 python ./run.py --depths 7 8 9

and then aggregate the results.

[helman44]

@merlinND Thanks so much for the insight on the ray tracing portion of Sionna. I'll more than likely end up using a BASH script similar to what you propose to do my splitting and concatenating. Thanks again!