Skip to content

Instantly share code, notes, and snippets.

@Akash-Rawat

Akash-Rawat/ModelSelection.py

Last active October 2, 2023 07:52
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save Akash-Rawat/700e7cc1284032ef8bd718f57e6df87b to your computer and use it in GitHub Desktop.
Save Akash-Rawat/700e7cc1284032ef8bd718f57e6df87b to your computer and use it in GitHub Desktop.
Download ZIP
Selecting a Pre-trained Model
Raw
ModelSelection.py
import torch
import torch.nn as nn
import torchvision.models as models
base_model = models.resnet18(pretrained=True)
num_Genders = len(image_datasets.classes)
num_ftrs = base_model.fc.in_features
base_model.fc = nn.Linear(num_ftrs, num_Genders)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment