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

29
.gitignore vendored
View File

@@ -1,27 +1,2 @@
# Python cache
__pycache__/
*.pyc
*.pyo
# Virtual environment
venv/
.venv/
bin/
lib/
lib64/
include/
share/
pyvenv.cfg
# Data & datasets
data/
bin/
# Model weights
*.pth
# System / logs
.DS_Store
*.log
*.tmp
# Created by venv; see https://docs.python.org/3/library/venv.html
*

View File

@@ -10,7 +10,7 @@ from architectures.Model import Model, Architecture
# numbre of classes
CLASS_SIZE = 20
# batch
BATCH_SIZE = 8
BATCH_SIZE = 16
# size of images per class trainset + testset
# 30 works best, more than that and we dont have enough data
@@ -33,7 +33,8 @@ RESOLUTION = 224
# - RESNET50
# - DENSENET121
# - INCEPTION
arch = Architecture.RESNET18
# - GOOGLENET
arch = Architecture.EFFICIENTNET
# DATA PREPARATION
# load data set and prepare
@@ -94,7 +95,7 @@ model.train(
# save.
torch.save(
model.get().state_dict(),
f'{arch.name}.pth')
f'{arch.name.lower()}.pth')
# done tuning
print('Model saved!')
@@ -118,4 +119,4 @@ print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
# Evaluate
model.evaluate(
loader = test_loader)
loader = test_loader)

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

View File

@@ -11,12 +11,12 @@ class ResNet18(Model):
m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
# freez all layers
for param in m.parameters():
param.requires_grad = False
#for param in m.parameters():
# param.requires_grad = False
# unfreez the last two
for param in m.layer3.parameters(): param.requires_grad = True
for param in m.layer4.parameters(): param.requires_grad = True
#for param in m.layer3.parameters(): param.requires_grad = True
#for param in m.layer4.parameters(): param.requires_grad = True
m.fc = nn.Linear(m.fc.in_features, self.size)
return m.to(self.device)