added two modfels

This commit is contained in:
2026-05-03 17:48:51 +02:00
parent a80d996b0c
commit 1c04344ad6
4 changed files with 23 additions and 37 deletions

View File

@@ -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()