added two modfels
This commit is contained in:
29
.gitignore
vendored
29
.gitignore
vendored
@@ -1,27 +1,2 @@
|
|||||||
|
# Created by venv; see https://docs.python.org/3/library/venv.html
|
||||||
# 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
|
|
||||||
|
|||||||
9
Tune.py
9
Tune.py
@@ -10,7 +10,7 @@ from architectures.Model import Model, Architecture
|
|||||||
# numbre of classes
|
# numbre of classes
|
||||||
CLASS_SIZE = 20
|
CLASS_SIZE = 20
|
||||||
# batch
|
# batch
|
||||||
BATCH_SIZE = 8
|
BATCH_SIZE = 16
|
||||||
|
|
||||||
# size of images per class trainset + testset
|
# size of images per class trainset + testset
|
||||||
# 30 works best, more than that and we dont have enough data
|
# 30 works best, more than that and we dont have enough data
|
||||||
@@ -33,7 +33,8 @@ RESOLUTION = 224
|
|||||||
# - RESNET50
|
# - RESNET50
|
||||||
# - DENSENET121
|
# - DENSENET121
|
||||||
# - INCEPTION
|
# - INCEPTION
|
||||||
arch = Architecture.RESNET18
|
# - GOOGLENET
|
||||||
|
arch = Architecture.EFFICIENTNET
|
||||||
|
|
||||||
# DATA PREPARATION
|
# DATA PREPARATION
|
||||||
# load data set and prepare
|
# load data set and prepare
|
||||||
@@ -94,7 +95,7 @@ model.train(
|
|||||||
# save.
|
# save.
|
||||||
torch.save(
|
torch.save(
|
||||||
model.get().state_dict(),
|
model.get().state_dict(),
|
||||||
f'{arch.name}.pth')
|
f'{arch.name.lower()}.pth')
|
||||||
|
|
||||||
# done tuning
|
# done tuning
|
||||||
print('Model saved!')
|
print('Model saved!')
|
||||||
@@ -118,4 +119,4 @@ print(f"Total test images for these {CLASS_SIZE} classes: {len(test_data)}")
|
|||||||
|
|
||||||
# Evaluate
|
# Evaluate
|
||||||
model.evaluate(
|
model.evaluate(
|
||||||
loader = test_loader)
|
loader = test_loader)
|
||||||
|
|||||||
@@ -1,3 +1,4 @@
|
|||||||
|
|
||||||
from abc import ABC, abstractmethod
|
from abc import ABC, abstractmethod
|
||||||
import torch
|
import torch
|
||||||
import torch.nn as nn
|
import torch.nn as nn
|
||||||
@@ -85,10 +86,17 @@ class Model(ABC):
|
|||||||
case Architecture.DENSENET121:
|
case Architecture.DENSENET121:
|
||||||
from architectures.DenseNet121 import DenseNet121
|
from architectures.DenseNet121 import DenseNet121
|
||||||
return DenseNet121(device, size)
|
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 _:
|
case _:
|
||||||
raise ValueError(f"Unknown model: {arch}")
|
raise ValueError(f"Unknown model: {arch}")
|
||||||
|
|
||||||
|
|
||||||
# model architectures
|
# model architectures
|
||||||
from enum import Enum, auto
|
from enum import Enum, auto
|
||||||
|
|
||||||
@@ -96,4 +104,6 @@ class Architecture(Enum):
|
|||||||
RESNET18 = auto()
|
RESNET18 = auto()
|
||||||
RESNET50 = auto()
|
RESNET50 = auto()
|
||||||
INCEPTION = auto()
|
INCEPTION = auto()
|
||||||
DENSENET121 = auto()
|
DENSENET121 = auto()
|
||||||
|
GOOGLENET = auto()
|
||||||
|
EFFICIENTNET = auto()
|
||||||
@@ -11,12 +11,12 @@ class ResNet18(Model):
|
|||||||
m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
|
m = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
|
||||||
|
|
||||||
# freez all layers
|
# freez all layers
|
||||||
for param in m.parameters():
|
#for param in m.parameters():
|
||||||
param.requires_grad = False
|
# param.requires_grad = False
|
||||||
|
|
||||||
# unfreez the last two
|
# unfreez the last two
|
||||||
for param in m.layer3.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
|
#for param in m.layer4.parameters(): param.requires_grad = True
|
||||||
|
|
||||||
m.fc = nn.Linear(m.fc.in_features, self.size)
|
m.fc = nn.Linear(m.fc.in_features, self.size)
|
||||||
return m.to(self.device)
|
return m.to(self.device)
|
||||||
Reference in New Issue
Block a user