File size: 1,321 Bytes
f33a809 |
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 |
import torch as T
from torch import nn, optim
import torch.nn.functional as F
import torchvision.models as models
from typing import Union, List
def get_model(num_classes:int, unfreeze_layers:Union[None, List[int]] = None, drop_rate: Union[None, float] = None):
model = models.efficientnet_b7(weights=models.EfficientNet_B7_Weights.DEFAULT)
for param in model.parameters():
param.requires_grad = False
if unfreeze_layers is not None and len(unfreeze_layers) > 0:
# Now unfreeze the layers in the unfreeze layer/ list
for layer_num in unfreeze_layers:
for name, child in model.features[layer_num].named_modules():
if not isinstance(child, nn.BatchNorm2d) and \
not isinstance(child, nn.Sequential) and \
not hasattr(child, 'block'):
for param in child.parameters():
param.requires_grad = True
if drop_rate is not None:
model.classifier[0] = nn.Dropout(drop_rate)
# Chagne the classifier head as per our need
model.classifier[1] = nn.Linear(2560, num_classes)
return model
if __name__ == "__main__":
model = get_model(
6,
[-1],
0.1
)
|