Spaces:
Runtime error
Runtime error
| 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) | |
| #) |