added two modfels
This commit is contained in:
@@ -1,3 +1,4 @@
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -85,10 +86,17 @@ class Model(ABC):
|
||||
case Architecture.DENSENET121:
|
||||
from architectures.DenseNet121 import DenseNet121
|
||||
return DenseNet121(device, size)
|
||||
# googleNet
|
||||
case Architecture.GOOGLENET:
|
||||
from architectures.GoogleNet import GoogleNet
|
||||
return GoogleNet(device, size)
|
||||
# EfficientNet
|
||||
case Architecture.EFFICIENTNET:
|
||||
from architectures.EfficentNet import EfficientNet
|
||||
return EfficientNet(device, size)
|
||||
case _:
|
||||
raise ValueError(f"Unknown model: {arch}")
|
||||
|
||||
|
||||
# model architectures
|
||||
from enum import Enum, auto
|
||||
|
||||
@@ -96,4 +104,6 @@ class Architecture(Enum):
|
||||
RESNET18 = auto()
|
||||
RESNET50 = auto()
|
||||
INCEPTION = auto()
|
||||
DENSENET121 = auto()
|
||||
DENSENET121 = auto()
|
||||
GOOGLENET = auto()
|
||||
EFFICIENTNET = auto()
|
||||
Reference in New Issue
Block a user