Pytorch - skip calculating features of pretrained models for every epoch
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, &quot;features_x.dat&quot;)  

# 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 &amp; train only the fully connected model
    
model.compile( loss=&quot;categorical_crossentropy&quot;, optimizer=keras.optimizers.Adam(learning_rate=0.001))
history = model.fit(  features_x, Y_train, batch_size=16, epochs=Epochs)


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

