[Python] InvalidArgumentError with tf.data.Dataset.from_generator() at random points during...

I want to use a python generator with a tf Dataset to extract and preprocess sequences from a (rather large) numpy array, but I keep getting this error at random points during the first epoch of training (anywhere between ~100 to ~1000 steps in). I'm using tensorflow 2.15.0

class BatchGenerator:
def __init__(
self,
scale_window: int = 3600,
depth: int = 50,
n_forecast: int = 30,
sequence_length: int = 1200,
polynomial_features_degree: int = 3,
batch_size: int = 1,
x_dtype: tf.dtypes.DType | str = "bfloat16",
y_dtype: tf.dtypes.DType | str = tf.float32,
):

self.scale_window = scale_window
self.depth = depth
self.n_forecast = n_forecast
self.sequence_length = sequence_length
self.batch_size = batch_size

self.epsilon = epsilon
self.x_dtype = x_dtype
self.y_dtype = y_dtype

# Initialize stateful preprocessing objects
self.mmx = MinMaxScaler()
self.poly = PolynomialFeatures(polynomial_features_degree, include_bias=False)

# Load CSV, format as np.ndarray
self.x_data = import_csv(config.filepath))
self.x_data = format_csv(self.x_data, depth)

# Detrend
self.price_data = self.x_data[1:, 0, [0, 2]]
self.price_data = tf.constant(self.price_data, y_dtype)

self.x_data[1:, :, [0, 2]] = np.subtract(
self.x_data[1:, :, [0, 2]], self.x_data[:-1, :, [0, 2]]
)
self.x_data[1:, :, [1, 3]] = np.subtract(
self.x_data[1:, :, [1, 3]], self.x_data[1:, :, [1, 3]]
)
self.x_data = self.x_data[1:]

self.indices = tf.range(
start=scale_window,
limit=len(self.x_data) - sequence_length - n_forecast + 1,
)
self.indices = tf.random.shuffle(self.indices, seed=1)

n_features = self.poly.fit_transform(np.zeros((depth, 2))).shape[-1]
self.x_batch_shape = (sequence_length, depth, n_features)

self.output_signature = (
tf.TensorSpec(
shape=(batch_size, *self.x_batch_shape),
dtype=x_dtype,
),
tf.TensorSpec(
shape=(batch_size, *self.x_batch_shape),
dtype=x_dtype,
),
tf.TensorSpec(
shape=(
batch_size,
2,
),
dtype=y_dtype,
),
tf.TensorSpec(shape=(batch_size, n_forecast, 2), dtype=self.y_dtype),
)

self.n_batches = len(self.indices) // batch_size
self.counter = 0

def generate(self):
while True:

batch_demand = tf.TensorArray(self.x_dtype, self.batch_size)
batch_supply = tf.TensorArray(self.x_dtype, self.batch_size)
batch_scale_params = tf.TensorArray(self.y_dtype, self.batch_size)
batch_targets = tf.TensorArray(self.y_dtype, self.batch_size)


for i in range(self.batch_size):
# Get X-values
x_start = self.indices[self.counter]

demand = self.x_data[
x_start - self.scale_window : x_start + self.sequence_length, :, :2
]
supply = self.x_data[
x_start - self.scale_window : x_start + self.sequence_length, :, 2:
]

# Generate interaction features
demand = self.poly.fit_transform(demand.reshape(-1, 2)).reshape(
(*demand.shape[:-1], -1)
)
supply = self.poly.fit_transform(supply.reshape(-1, 2)).reshape(
(*supply.shape[:-1], -1)
)

# Scale from last to first sample
demand = np.flip(demand, 0)
supply = np.flip(supply, 0)

# Scale
for j in range(self.sequence_length):
for k in range(demand.shape[-1]):
self.mmx.fit(
np.concatenate(
(
demand[j : j + self.scale_window + 1, :, k],
supply[j : j + self.scale_window + 1, :, k],
),
axis=0,
).reshape(-1, 1)
)
demand[j, :, k] = self.mmx.transform(
demand[j, :, k].reshape(-1, 1)
).reshape(demand[j, :, k].shape)

supply[j, :, k] = self.mmx.transform(
supply[j, :, k].reshape(-1, 1)
).reshape(supply[j, :, k].shape)

if (j == 0) and (k == 0):
batch_scale_params = batch_scale_params.write(
i,
tf.constant(
[self.mmx.data_min_[0], self.mmx.data_max_[0]],
self.y_dtype,
),
)

demand = demand[: self.sequence_length]
supply = supply[: self.sequence_length]
demand = np.flip(demand, 0)
supply = np.flip(supply, 0)


batch_demand = batch_demand.write(i, tf.constant(demand, self.x_dtype))

batch_supply = batch_supply.write(i, tf.constant(supply, self.x_dtype))

# Get targets
forecast_start = x_start + self.sequence_length

batch_targets = batch_targets.write(
i,
tf.constant(
self.x_data[
forecast_start : forecast_start + self.n_forecast,
0,
[0, 2],
],
self.y_dtype,
),
)

# Update counter
if self.counter + 1 == len(self.indices):
self.counter = 0
else:
self.counter += 1

batch_demand = batch_demand.stack()
batch_supply = batch_supply.stack()
batch_scale_params = batch_scale_params.stack()
batch_targets = batch_targets.stack()

yield (batch_demand, batch_supply, batch_scale_params, batch_targets)


Traceback (most recent call last):
File "/Users/.../Desktop/virtual_envs/scripts_keras2/varformer_train.py", line 108, in <module>
for i, batch in enumerate(data):
File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 810, in __next__
return self._next_internal()
File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 773, in _next_internal
ret = gen_dataset_ops.iterator_get_next(
File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 3029, in iterator_get_next
_ops.raise_from_not_ok_status(e, name)
File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/framework/ops.py", line 5883, in raise_from_not_ok_status
raise core._status_to_exception(e) from None # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.UnknownError: {{function_node __wrapped__IteratorGetNext_output_types_4_device_/job:localhost/replica:0/task:0/device:CPU:0}} InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Index out of range using input dim 1; input has only 1 dims [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/ops/script_ops.py", line 270, in __call__
ret = func(*args)

File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
return func(*args, **kwargs)

File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/data/ops/from_generator_op.py", line 198, in generator_py_func
values = next(generator_state.get_iterator(iterator_id))

File "/Users/.../Desktop/virtual_envs/scripts_keras2/_generator.py", line 500, in generate
x_start = self.indices[self.counter]

File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
raise e.with_traceback(filtered_tb) from None

File "/Users/.../Desktop/virtual_envs/scripts_keras2/.venv/lib/python3.10/site-packages/tensorflow/python/framework/ops.py", line 5883, in raise_from_not_ok_status
raise core._status_to_exception(e) from None # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Index out of range using input dim 1; input has only 1 dims [Op:StridedSlice] name: strided_slice/


[[{{node PyFunc}}]] [Op:IteratorGetNext] name:


If I iterate over the tf dataset, everything works as expected (counter resets correctly, each batch has the same dimensions, etc...). I thought it might be an issue with my model, but training with the standalone generator works fine, as well (besides taking 3 times longer without being able to use prefetch()).

I've also tried using building batches in generate() using tensor_scatter_nd_update() or by appending samples to lists, then stacking them to see if this is a problem with TensorArrays, but run into the same problem, there.

Continue reading...