Edit: As you increase the size of the model, the difference in speeds disappears. This leads me to believe the difference is in memory management. As the network gets larger the percent of time it's spending loading the next batch gets lower. With the tiny network in my example the memory allocation is actually taking up a good chunk of the total time because the forward/backward passes are almost instant. Fit() can know about all the data at once and just iterate a pointer across it (or something?)
Training an epoch using model.fit is 2-3x faster than looping over your data with model.train_on_batch. I don't understand how because I've always been told both do the same thing. GPU usage is the same with either method so I can't imagine where all that extra time is being spent.
Here is minimal code to test each version.
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.layers import Input, Conv2D, Flatten, Dense
import numpy as np
import time
import tensorflow as tf
def makeNet():
img = Input(shape=(28, 28, 1))
x = Conv2D(128, 3, activation="relu", strides=2, padding="same")(img)
x = Conv2D(64, 3, activation="relu", strides=2, padding="same")(x)
x = Conv2D(32, 3, activation="relu", strides=2, padding="same")(x)
x = Flatten()(x)
x = Dense(10, activation="sigmoid")(x)
model = Model(img, x)
opt = Adam()
model.compile(loss='mse', optimizer=opt)
return model
(digits, labels), (_, _) = tf.keras.datasets.mnist.load_data()
labels = labels.astype("float32") / 10
digits = np.expand_dims(digits, -1).astype("float32") / 255
bs = 128
model = makeNet()
while True:
st = time.perf_counter()
model.fit(digits, labels, batch_size = bs, epochs=1)
# remove the above line and uncomment lines below to test the train on batch version
# for i in range(len(digits) // bs):
# x = digits[i * bs : i * bs + bs]
# y = labels[i * bs : i * bs + bs]
#model.train_on_batch(x, y)
print('time =', time.perf_counter() - st)
I have tried two different models with the same result. I have tried looking at the code for fit() but the way it's written is really hard to follow and I can't tell what's going on.
The problem with this is that I always use train on batch because my datasets never fit in memory and I also like to visualize results at a higher interval than just once per epoch.
Is it possible to train on individual batches at the same speed as fit()? What is making fit() so much faster?