model.py

import torch
import torchvision

from torch import nn

def create_effnetb1_model(num_classes:int=101):
  # 1,2,3 EffNetB1의 미리 훈련된 weights, transforms, model 얻기
  weights = torchvision.models.EfficientNet_B1_Weights.DEFAULT
  transforms = weights.transforms()
  model = torchvision.models.efficientnet_b1(weights=weights)

  # 4. 기본 모델의 layer 고정
  for param in model.parameters():
    param.requires_grad = False
  # 5. classifier head 바꾸기
  model.classifier = nn.Sequential(
      nn.Dropout(p=0.3, inplace=True),
      nn.Linear(in_features=1280, out_features=num_classes)
  )
  return model, transforms