Pytorch - skip calculating features of pretrained models for every epoch - keras

Question

I am used to work with tenserflow - keras but now I am forced to start working with Pytorch for flexibility issues. However, I don't seem to find a pytorch code that is focused on training only the classifciation layer of a model. Is that not a common practice ? Now I have to wait out the calculation of the feature extraction of the same data for every epoch. Is there a way to avoid that ?
# in tensorflow - keras :
from tensorflow.keras.applications import vgg16, MobileNetV2, mobilenet_v2
# Load a pre-trained
pretrained_nn = MobileNetV2(weights='imagenet', include_top=False, input_shape=(Image_size, Image_size, 3))
# Extract features of the training data only once
X = mobilenet_v2.preprocess_input(X)
features_x = pretrained_nn.predict(X)
# Save features for later use
joblib.dump(features_x, "features_x.dat")
# Create a model and add layers
model = Sequential()
model.add(Flatten(input_shape=features_x.shape[1:]))
model.add(Dense(100, activation='relu', use_bias=True))
model.add(Dense(Y.shape[1], activation='softmax', use_bias=False))
# Compile & train only the fully connected model
model.compile( loss="categorical_crossentropy", optimizer=keras.optimizers.Adam(learning_rate=0.001))
history = model.fit( features_x, Y_train, batch_size=16, epochs=Epochs)

Answer 1

Assuming you already have the features ìn features_x, you can do something like this to create and train the model:
# create a loader for the data
dataset = torch.utils.data.TensorDataset(features_x, Y_train)
loader = torch.utils.data.DataLoader(dataset, batch_size=16, shuffle=True)
# define the classification model
in_features = features_x.flatten(1).size(1)
model = torch.nn.Sequential(
torch.nn.Flatten(),
torch.nn.Linear(in_features=in_features, out_features=100, bias=True),
torch.nn.ReLU(),
torch.nn.Linear(in_features=100, out_features=Y.shape[1], bias=False) # Softmax is handled by CrossEntropyLoss below
)
model.train()
# define the optimizer and loss function
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_function = torch.nn.CrossEntropyLoss()
# training loop
for e in range(Epochs):
for batch_x, batch_y in enumerate(loader):
optimizer.zero_grad() # clear gradients from previous batch
out = model(batch_x) # forward pass
loss = loss_function(out, batch_y) # compute loss
loss.backward() # backpropagate, get gradients
optimizer.step() # update model weights