File size: 1,137 Bytes
5f37d56
 
 
 
 
 
 
 
08f8524
5f37d56
 
 
 
 
 
 
 
 
 
e315722
5f37d56
 
 
 
 
 
 
 
 
3119e39
5f37d56
 
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
from typing import Tuple

import torch
from torch import nn
import torchvision

def create_effnetb2_model(num_classes: int = 3,
                          seed: int = 4,
                         ) -> Tuple[nn.Module, torchvision.transforms.Compose]:
    """Create an EfficientNetB2 feature extractor model and transforms.
    
    Args:
      num_classes: Number of classes to use for classification (default 3).
      seed: Random seed for reproducibility (default 4).
      
    Returns:
      A tuple (model, transforms) of the model and its image transforms.
    """
    weights = torchvision.models.EfficientNet_B2_Weights.DEFAULT
    transforms = weights.transforms()
    model = torchvision.models.efficientnet_b2(weights=weights)
    
    # Freeze parameters below the head
    for param in model.parameters():
        param.requires_grad = False
    # Replace the classifier head with one of appropriate size for the problem
    torch.manual_seed(seed)
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.3, inplace=True),
        nn.Linear(in_features=1408, out_features=num_classes)
    )
    return model, transforms