I was trying to make a model testing series I ran into some horrible compilation behavior. The model runs about 5x faster on my CPU (Xeon W-2225, 4 core) than on my RTX 3060 (12GB) OR on my T1000 (4GB). During the run, during compilation I get hundreds of messages like
... local_xla/xla/stream_executor/cuda/subprocess_compilation.cc:346] ptxas warning : Registers are spilled to local memory in function 'gemm_fusion_dot_98'
HUNDREDS and thousands. Looking at the TensorBoard profiler I can see that 99.9% of tasks are running on 'Host' (I think that means CPU). And all of the step time is taken up by 'All other' category.
More over, I have mixed_precision.Policy('mixed_float16') enabled, but 100% of computation was in float32. (As far as I can tell I followed the TF tutorial on MixedPrec.)
The model is fairly small, but there is clearly something wrong. Why would I get so many ptxas messages?
Just so you can read through it and maybe recreate the issue, here's my setup and code:
SETUP:
- OS Ubuntu
- GPU RTX 3060 12GB,
- Driver: 575.51.03 (installed through Ubuntu's Additional Drivers app)
- CUDA: 12.9
- CPU Xeon W-2225
- TensorFlow 2.19, installed via pip (with cuda as per TF install guide)
CODE It is probably a bit silly looking, but 1. I'm a noob and 2nd I tried to make it easy to do Hyperparameter searches on multiple models all together, so I can let it run overnight and evaluate in the morning.
import numpy as np
import keras
import tensorflow as tf
import matplotlib.pyplot as plt
import keras_tuner as kt
from tensorflow.keras import layers
from sionna_data_loader import *
from IPython.display import clear_output
from tensorflow.keras import mixed_precision
from datetime import datetime
import os
clear_output() #clear warnings of not using AVX2, AVX512, FMA and so and so as TF was installed with pip
os.environ["CUDA_VISIBLE_DEVICES"] = "0" #force to use RTX3060
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)
Data loading here random as example, but the actual data looks and loads in a similar way, just from .npz files (read into memory and built to Dataset with .from_tensor_slices()).
CLASSES = 72
VALIDATION_SPLIT = 0.2
TEST_SPLIT = 0.125
ZERO_MEAN_UNIT_DEV = True
raw_data = np.random.random(size=(20000, 4)).astype('float32')
raw_labels = np.random.randint(size=(20000,), low=0, high=360).astype('int16')
binned_labels = np.floor(raw_labels/360*CLASSES).astype('int16') #bin to classes
dat_len = raw_data.shape[0]
X_train = raw_data[ :int(dat_len*(1-VALIDATION_SPLIT-TEST_SPLIT))] #select training array
y_train = binned_labels[ :int(dat_len*(1-VALIDATION_SPLIT-TEST_SPLIT))] #select training labels
X_val = raw_data[int(dat_len*(1-VALIDATION_SPLIT-TEST_SPLIT)) : int(dat_len*(1-TEST_SPLIT))]
y_val = binned_labels[int(dat_len*(1-VALIDATION_SPLIT-TEST_SPLIT)) : int(dat_len*(1-TEST_SPLIT))]
X_test = raw_data[int(dat_len*(1-TEST_SPLIT)): ]
y_test = binned_labels[int(dat_len*(1-TEST_SPLIT)): ]
if ZERO_MEAN_UNIT_DEV:
train_stdev = np.std(X_train)
train_mean = np.mean(X_train)
X_train = (X_train-train_mean)/train_stdev
X_val = (X_val-train_mean)/train_stdev
X_test = (X_test-train_mean)/train_stdev
train_ds = tf.data.Dataset.from_tensor_slices((X_train, y_train)).batch(np.floor(X_train.shape[0]/10)).prefetch(tf.data.AUTOTUNE)
val_ds = tf.data.Dataset.from_tensor_slices((X_val, y_val)).batch(np.floor(X_val.shape[0]/10)).prefetch(tf.data.AUTOTUNE)
test_ds = tf.data.Dataset.from_tensor_slices((X_test, y_test)).batch(np.floor(X_test.shape[0]/10)).prefetch(tf.data.AUTOTUNE)
del raw_data
del raw_labels
del binned_labels
del X_train
del X_val
del X_test
del y_train
del y_val
del y_test
Helper function for easy model HP searches:
def search_hyperModel(name, model:keras.Model):
tuner = kt.Hyperband(
model,
objective='val_accuracy',
max_epochs=10,
factor=3,
directory='hyperModels',
project_name=name)
stop_early = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=4)
tuner.search(
train_ds,
validation_data=val_ds,
epochs=50,
callbacks=[stop_early],
)
# Get the optimal hyperparameters
best_hps=tuner.get_best_hyperparameters(num_trials=1)[0]
print(f"The hyperparameter search is complete.")
model = tuner.hypermodel.build(best_hps)
best_only_callback = keras.callbacks.ModelCheckpoint(
f"/home/path/hyperModels/{name}/{name}.keras",
monitor='val_loss',
verbose=0,
save_best_only=True,
save_weights_only=False,
mode='auto',
save_freq='epoch',
)
early_stop_callback = keras.callbacks.EarlyStopping(
monitor='val_accuracy',
patience=5,
restore_best_weights=True
)
reduce_lr_callback = keras.callbacks.ReduceLROnPlateau(
monitor='val_loss',
factor=0.1,
patience=2,
min_lr=0.00001
)
tboard_callback = tf.keras.callbacks.TensorBoard(
log_dir = f"/home/path/logs/{name}",
histogram_freq = 1,
profile_batch = '10,12'
)
history = model.fit(
train_ds,
validation_data=val_ds,
epochs=50,
callbacks=[early_stop_callback, best_only_callback, reduce_lr_callback, tboard_callback],
verbose=0
)
fig, axs = plt.subplots(1,2)
axs[0].plot(history.history['accuracy'], label="train acc%")
axs[0].plot(history.history['val_accuracy'], '--', label="val acc%")
axs[0].grid()
axs[0].legend()
axs[1].plot(history.history['loss'], label="train loss")
axs[1].plot(history.history['val_loss'], '--', label="val loss")
axs[1].grid()
axs[1].legend()
fig.set_size_inches(14,5)
plt.show()
model = keras.models.load_model(f"/home/path/hyperModels/{name}/{name}.keras")
loss, acc = model.evaluate(test_ds, verbose=2)
print(f"Optimized model acc: {acc*100:.1f}% after {len(history.history['accuracy'])} epochs.")
return {'model': model, 'name':name}
Build model, as simple as it gets:
def simple_dense(hp):
hp_units = hp.Int('units', min_value=8, max_value=512, step=8)
hp_activ = hp.Choice('activation_f', ['relu', 'softplus', 'celu', 'glu', 'tanh', 'linear', 'mish', 'sigmoid', 'silu', 'softsign'])
model = keras.Sequential(layers=[
layers.Input(shape=(4,)),
layers.Dense(hp_units, activation=hp_activ),
layers.Dense(CLASSES, activation='softmax', dtype='float32'),
])
# Tune the learning rate for the optimizer
hp_learning_rate = hp.Choice('learning_rate', values=np.logspace(-5, -2, base=10, num=20).tolist())
model.compile(optimizer=keras.optimizers.Adam(learning_rate=hp_learning_rate),
loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
metrics=['accuracy'])
return model
hpModel = search_hyperModel(
name=f"simple_dense_{datetime.now().strftime(r'%m-%d_%H-%M')}",
model = simple_dense
)
models = {}
models[hpModel['name']] = hpModel['model']
Finally
%tensorboard
As I said, the code runs much faster on CPU only due to the many warning messages....