how to overfit a model on a single batch in keras?

Question

I am trying to overfit my model on a single batch to check model integrity. I am using Keras and TensorFlow for the implementation of my model and coding style for this project.

I know how to get the single batch and overfit the model in PyTorch but don't have an idea in Keras.

to get a single batch in PyTorch I used:

images, labels = next(iter(train_dataset))
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr = 0.0001)

for epoch in range(epochs):
    print(f"Epoch [{epoch}/{epochs}]")
    # for batch_idx, (data, target) in enumerate(train_loader):
    data, target = data.to(device), target.to(device)
    data = data.reshape(data.shape[0], -1)
    # forward
    score = model(data)
    loss = criterion(score, target)
    print(f"Loss: {loss.item()}")
        # backward
    optimizer.zero_grad()
    loss.backward()
    
    optimizer.step() 

How to do it in keras any helping matrial?

Answer 1

Thank you everyone for coming here. I found a solution and here it is:

datagen = ImageDataGenerator(rescale=1 / 255.0,
                            rotation_range=20,
                            zoom_range=0.2,
                            width_shift_range=0.05,
                            height_shift_range=0.05,
                            shear_range=0.2,
                            horizontal_flip=True,
                            fill_mode="nearest"
                            )
# preprocessing_function=preprocess_input,
# Declare an image generator for validation & testing without generation
test_datagen = ImageDataGenerator(rescale = 1./255,)#preprocessing_function=preprocess_input

# Declare generators for training, validation, and testing from DataFrames
train_gen = datagen.flow_from_directory(directory_train,
                                        target_size=(512, 512),
                                        color_mode='rgb',
                                        batch_size=BATCH_SIZE,
                                        class_mode='binary',
                                        shuffle=True)

val_gen = test_datagen.flow_from_directory(directory_val,
                                        target_size=(512, 512),
                                        color_mode='rgb',
                                        batch_size=BATCH_SIZE,
                                        class_mode='binary',
                                        shuffle=False)

test_gen = test_datagen.flow_from_directory(directory_test,
                                        target_size=(512, 512),
                                        color_mode='rgb',
                                        batch_size=BATCH_SIZE,
                                        class_mode='binary',
                                        shuffle=False)
train_images, train_labels = next(iter(train_gen))
val_images, val_labels = next(iter(val_gen))
test_images, test_labels = next(iter(val_gen))

#check shape for selected Batch

print("Length of Train images : {}".format(len(train_images)))
print("shape of Train images : {}".format(train_images.shape))
print("shape of Train labels : {}".format(train_labels.shape))

Length of Train images : 32
shape of Train images : (32, 512, 512, 3)
shape of Train labels : (32,)

history = model.fit(train_images, train_labels,
                            use_multiprocessing=True,
                            workers=16,
                            epochs=100,     
                            class_weight=class_weights,                 
                            validation_data=(val_images, val_labels),                    
                            shuffle=True,
                            callbacks=call_backs)