Access individual gradients - TensorFlow2 - python-3.x

Question

For a toy LeNet-5 CNN architecture on MNIST implemented in TensorFlow-2.10 + Python-3.10, with a batch-size = 256:
class LeNet5(Model):
def __init__(self):
super(LeNet5, self).__init__()
self.conv1 = Conv2D(
filters = 6, kernel_size = (5, 5),
strides = (1, 1), activation = None,
input_shape = (28, 28, 1)
)
self.pool1 = AveragePooling2D(
pool_size = (2, 2), strides = (2, 2)
)
self.conv2 = Conv2D(
filters = 16, kernel_size = (5, 5),
strides = (1, 1), activation = None
)
self.pool2 = AveragePooling2D(
pool_size = (2, 2), strides = (2, 2)
)
self.flatten = Flatten()
self.dense1 = Dense(
units = 120, activation = None
)
self.dense2 = Dense(
units = 84, activation = None
)
self.output_layer = Dense(
units = 10, activation = None
)
def call(self, x):
x = tf.nn.relu(self.conv1(x))
x = self.pool1(x)
x = tf.nn.relu(self.conv2(x))
x = self.pool2(x)
x = self.flatten(x)
x = tf.nn.relu(self.dense1(x))
x = tf.nn.relu(self.dense2(x))
x = tf.nn.softmax(self.output_layer(x))
return x
def shape_computation(self, x):
print(f"Input shape: {x.shape}")
x = self.conv1(x)
print(f"conv1 output shape: {x.shape}")
x = self.pool1(x)
print(f"pool1 output shape: {x.shape}")
x = self.conv2(x)
print(f"conv2 output shape: {x.shape}")
x = self.pool2(x)
print(f"pool2 output shape: {x.shape}")
x = self.flatten(x)
print(f"flattened shape: {x.shape}")
x = self.dense1(x)
print(f"dense1 output shape: {x.shape}")
x = self.dense2(x)
print(f"dense2 output shape: {x.shape}")
x = self.output_layer(x)
print(f"output shape: {x.shape}")
del x
return None
# Initialize an instance of LeNet-5 CNN-
model = LeNet5()
model.build(input_shape = (None, 28, 28, 1))
# Define loss and optimizer-
loss_fn = tf.keras.losses.CategoricalCrossentropy(reduction = tf.keras.losses.Reduction.NONE)
# optimizer = tf.keras.optimizers.Adam(learning_rate = 0.0003)
optimizer = tf.keras.optimizers.SGD(
learning_rate = 10e-3, momentum = 0.0,
nesterov = False
)
with tf.GradientTape() as grad_tape:
pred = model(x)
loss = loss_fn(y, pred)
loss.shape
TensorShape([256])
This computes individual loss for each of the 256 training images in a given batch.
# Compute gradient using loss wrt parameters-
grads = grad_tape.gradient(loss, model.trainable_variables)
type(grads), len(grads)
# (list, 10)
for i in range(len(grads)):
print(f"i: {i}, grads.shape: {grads[i].shape}")
"""
i: 0, grads.shape: (5, 5, 1, 6)
i: 1, grads.shape: (6,)
i: 2, grads.shape: (5, 5, 6, 16)
i: 3, grads.shape: (16,)
i: 4, grads.shape: (256, 120)
i: 5, grads.shape: (120,)
i: 6, grads.shape: (120, 84)
i: 7, grads.shape: (84,)
i: 8, grads.shape: (84, 10)
i: 9, grads.shape: (10,)
"""
Corresponding to loss for each training example, how can I compute gradient corresponding to each training example?