File size: 4,178 Bytes
74c02ff |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import torch
import torch.nn as nn
import torchvision.models as models
from torchvision.models import VGG16_Weights
class VGG16Model(nn.Module):
def __init__(self, num_classes=4):
super(VGG16Model, self).__init__()
#self.vgg16 = models.vgg16(pretrained=True)
# Replace pretrained=True with weights=VGG16_Weights.DEFAULT
#self.vgg16 = models.vgg16(weights=VGG16_Weights.DEFAULT)
# Manually load the weights
self.vgg16 = models.vgg16()
# state_dict = torch.load("/work/mech-ai/angona3/Trial/vgg16_weights.pth", map_location=torch.device("cpu"))
# self.vgg16.load_state_dict(state_dict)
self.vgg16.classifier[6] = nn.Linear(4096, num_classes) ##In PyTorch CrossEntropyLoss Handles Softmax Internally
###Adding an explicit softmax layer would create numerical instability by double-applying softmax (once in your model, once in the loss),
# leading to incorrect gradient calculations during training.
self.vgg16.classifier[2] = nn.Dropout(p=0.6) # or 0.7 ##introducing dropout for regularization
self.vgg16.classifier[5] = nn.Dropout(p=0.6) # or 0.7
# Freeze all the convolutional layers (feature extractor part)
# The classifier layers (fully connected layers) remain trainable
for param in self.vgg16.features.parameters():
param.requires_grad = False
#param.requires_grad = True # Unfreeze the last two fully connected layers
# Unfreeze Conv Block 4 and Conv Block 5 (512 filters, 3x3 filters, same padding)
# best: conv_layers_to_unfreeze = [17, 19, 21, 24, 26, 28]
conv_layers_to_unfreeze = [17, 19, 21, 24, 26, 28]
for layer_idx in conv_layers_to_unfreeze:
for param in self.vgg16.features[layer_idx].parameters():
param.requires_grad = True
# Unfreeze the last two fully connected layers
# Unfreeze all fully connected layers
for param in self.vgg16.classifier.parameters():
param.requires_grad = True
def forward(self, x):
return self.vgg16(x)
#Fine-tuning the entire network can lead to better performance compared to freezing layers
##because the model can adjust both the feature extractor and classifier to your specific dataset.
##Yes, nn.CrossEntropyLoss in PyTorch is explicitly designed to handle raw logits (ranging from −∞to +∞) directly and efficiently.
##CrossEntropyLoss=Softmax(logits)+Log+NLLLoss
##### In newer versions of torchvision, the weights argument replaces pretrained.
# #The VGG16_Weights.DEFAULT is the new way to specify that anyone want the pretrained weights, and it is the preferred method.
##Deprecation Warning Fix: Replace pretrained=True with weights=VGG16_Weights.DEFAULT.
##Security Warning Fix: You don't need to use torch.load in this case because you're not loading a pre-trained model from a .pth file.
##from torchvision.models import VGG16_Weights
## self.vgg16 = models.vgg16(weights=VGG16_Weights.DEFAULT)
## adding dropout to VGG16
#class VGG16Model(nn.Module):
# def __init__(self, num_classes=4):
# super(VGG16Model, self).__init__()
# self.vgg16 = models.vgg16(pretrained=True)
# self.vgg16.classifier[6] = nn.Sequential(
# nn.Linear(4096, 1024), # First layer reduces dimensions to 1024
# nn.ReLU(), # Adds non-linearity to increase learning capacity
# nn.Dropout(0.5), # Introduces dropout for regularization
# nn.Linear(1024, num_classes) # Second layer maps to the desired number of classes
# )
# def forward(self, x):
# return self.vgg16(x)
#import torch.nn as nn
#from torchvision import models
#def get_resnet18_model(num_classes):
# model = models.resnet18(pretrained=True)
# model.fc = nn.Linear(model.fc.in_features, num_classes)
# return model
##VGG16 structure
#(vgg16.classifier): Sequential(
# (0): Linear(25088, 4096)
# (1): ReLU(inplace=True)
# (2): Dropout(p=0.5, inplace=False)
# (3): Linear(4096, 4096)
# (4): ReLU(inplace=True)
# (5): Dropout(p=0.5, inplace=False)
# (6): Linear(4096, num_classes)
#) |