aaa3k5/FoodVision_big

1import torch 2import torchvision3 4from torch import nn 5 6def create_effnetb2_model(num_classes:int =3,7                          seed:int= 42):8  """Create an EfficientNetB2 feature extractor model and transforms.9 10  Args:11      num_classes (int, optional): number of classes in the classifier head.12          Defaults to 3.13      seed (int, optional): random seed value. Defaults to 42.14  15  Returns:16      model (torch.nn.Module): EffNetB2 feature extractor model.17      transforms (torchvision.transforms): EffNetB2 image transforms.18  """19  weights= torchvision.models.EfficientNet_B2_Weights.DEFAULT20  transforms= weights.transforms()21  model= torchvision.models.efficientnet_b2(weights= weights)22 23  for param in model.parameters():24    param.requires_grad= False25  26  torch.manual_seed(seed)27  model.classifier= nn.Sequential(28      nn.Dropout(p= 0.3, inplace= True),29      nn.Linear(in_features= 1408, out_features= num_classes)30  )31 32  return model, transforms33