Getting the flattened image from a pretrained model

vision
1

How do I get a flattened image from a pretrained model (VGG16). I need to put this flattened image as an input to a simple perceptron and then train. Kindly help me with this part.

#Function to get pre-trained model
from torchvision import models
import torch.nn as nn
def get_pretrained_model(model_name):
    """Retrieve a pre-trained model from torchvision

    Params
    -------
        model_name (str): name of the model (currently only accepts vgg16 and resnet50)

    Return
    --------
        model (PyTorch model): cnn

    """

    if model_name == 'vgg16':
        model = models.vgg16(pretrained=True)

        # Freeze early layers
        for param in model.parameters():
            param.requires_grad = False
        n_inputs = model.classifier[6].in_features

        # Add on classifier
        model.classifier[6] = nn.Sequential(
            nn.Linear(n_inputs, 512), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(512, 5))

    elif model_name == 'resnet50':
        model = models.resnet50(pretrained=True)

        for param in model.parameters():
            param.requires_grad = False

        n_inputs = model.fc.in_features
        model.fc = nn.Sequential(
            nn.Linear(n_inputs, 512), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(512, 5))
    return model

#getting Pre-trained model
    model = models.vgg16(pretrained=True)
    # Freeze model weights
    for param in model.parameters():
        param.requires_grad = False
    
    n_inputs = model.classifier[6].in_features

    # Add on classifier
    model.classifier[6] = nn.Sequential(
        nn.Linear(n_inputs, 512), nn.ReLU(), nn.Dropout(0.4),
        nn.Linear(512, 5))

    model = get_pretrained_model('vgg16')
    print(model.forward)
    criterion = torch.nn.CrossEntropyLoss()
    n_class=5
    optimizer = torch.optim.Adam(model.parameters())
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min' if n_class > 1 else 'max', patience=2)
2

You could use the flatten operation of tensors!

input = torch.ones((1,3,100,100))
flat_input = input.flatten(input, start_dim=1)

Keep an eye on the start_dim parameter, so you keep the number of batches.