refactor api
This commit is contained in:
@@ -0,0 +1,41 @@
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass
|
||||
class Config:
|
||||
"""Settings for the YOLO face detector. Paths are relative to api/.
|
||||
Every field can be overridden on the command line, e.g. --batch-size 32."""
|
||||
|
||||
# Data (COCO format: images + _annotations.coco.json in each folder)
|
||||
train_dir: str = "data/faces_scenery/train"
|
||||
val_dir: str = "data/faces_scenery/test"
|
||||
annotation_file: str = "_annotations.coco.json"
|
||||
augment_train: bool = True
|
||||
augment_val: bool = True # the original code augmented validation too
|
||||
image_size: int = 64
|
||||
grid: int = 6
|
||||
|
||||
# Model
|
||||
c_hidden: int = 16
|
||||
boxes: int = 1
|
||||
labels: int = 1
|
||||
size_activation: str = "sigmoid" # "sigmoid" (current training) or "exp" (deployed model)
|
||||
|
||||
# Training
|
||||
epochs: int = 500
|
||||
batch_size: int = 64
|
||||
learning_rate: float = 1e-3
|
||||
weight_decay: float = 1e-4
|
||||
train_workers: int = 2
|
||||
val_workers: int = 0
|
||||
device: str = "auto" # "auto", "cpu", "cuda:0", ...
|
||||
|
||||
# Output
|
||||
save_dir: str = "saved_models"
|
||||
model_name: str = "face_detection_yolo"
|
||||
save: bool = True
|
||||
|
||||
# sample / webcam tasks
|
||||
checkpoint: str = "" # empty = <save_dir>/<model_name>/<model_name>
|
||||
threshold: float = 0.95 # confidence needed to draw a box
|
||||
camera: int = 0 # webcam index
|
||||
@@ -1,24 +1,14 @@
|
||||
import math
|
||||
import os
|
||||
import cv2
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from pycocotools.coco import COCO
|
||||
import numpy as np
|
||||
from torchvision import transforms
|
||||
import albumentations as A
|
||||
|
||||
from src.common.boxes import image_to_cell_coords, is_center_in_grid_cell
|
||||
|
||||
def is_center_in_grid_cell(x, y, img_w, img_h, S, cell_i, cell_j):
|
||||
cell_w = img_w / S
|
||||
cell_h = img_h / S
|
||||
|
||||
# find which cell the center belongs to
|
||||
gt_cell_i = min(int(x / cell_w), S - 1)
|
||||
gt_cell_j = min(int(y / cell_h), S - 1)
|
||||
|
||||
return (gt_cell_i == cell_i) and (gt_cell_j == cell_j)
|
||||
|
||||
def one_hot(index, num_classes):
|
||||
if index < 0 or index >= num_classes:
|
||||
@@ -28,18 +18,6 @@ def one_hot(index, num_classes):
|
||||
encoding[index] = 1
|
||||
return encoding
|
||||
|
||||
def turn_grid_centered(x, y, img_w, img_h, S, cell_i, cell_j):
|
||||
cell_w = img_w / S
|
||||
cell_h = img_h / S
|
||||
|
||||
cell_border_w = cell_w * cell_i
|
||||
cell_border_h = cell_h * cell_j
|
||||
|
||||
return x - cell_border_w, y - cell_border_h
|
||||
|
||||
def flip_bbox_horizontal(box, image_width):
|
||||
x, y, w, h = box
|
||||
return (image_width - x - w, y, w, h)
|
||||
|
||||
class YoloDataset(Dataset):
|
||||
def __init__(self, image_dir, annotation_path, img_size=64, grid = 9, transform=False):
|
||||
@@ -135,7 +113,7 @@ class YoloDataset(Dataset):
|
||||
img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y):
|
||||
|
||||
class_one_hot = one_hot(int(label), self.num_classes)
|
||||
x_grid_centered, y_grid_centered = turn_grid_centered(x=x_center, y=y_center, img_w=self.img_size,
|
||||
x_grid_centered, y_grid_centered = image_to_cell_coords(x=x_center, y=y_center, img_w=self.img_size,
|
||||
img_h=self.img_size, S=self.grid, cell_i=x, cell_j=y)
|
||||
|
||||
|
||||
@@ -0,0 +1,149 @@
|
||||
"""YOLO face detector tasks. Run with: python -m src.main yolo.<task>"""
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision.utils import draw_bounding_boxes
|
||||
|
||||
from src.common.checkpoints import checkpoint_path, load_checkpoint
|
||||
from src.common.device import get_device
|
||||
from src.common.training import benchmark_num_workers, fit
|
||||
from src.common.visualize import show_image
|
||||
from src.models.yolo import Yolo_model, convert_prediction
|
||||
from src.yolo.config import Config
|
||||
from src.yolo.dataset import YoloDataset
|
||||
from src.yolo.loss import YoloLoss
|
||||
|
||||
|
||||
def build_dataset(cfg: Config, folder: str, augment: bool):
|
||||
return YoloDataset(
|
||||
image_dir=folder,
|
||||
annotation_path=os.path.join(folder, cfg.annotation_file),
|
||||
img_size=cfg.image_size,
|
||||
transform=augment,
|
||||
grid=cfg.grid,
|
||||
)
|
||||
|
||||
|
||||
def build_model(cfg: Config):
|
||||
return Yolo_model(c_in=3, boxes=cfg.boxes, grid=cfg.grid, labels=cfg.labels,
|
||||
c_hidden=cfg.c_hidden, size_activation=cfg.size_activation)
|
||||
|
||||
|
||||
def load_trained(cfg: Config, device):
|
||||
path = cfg.checkpoint or checkpoint_path(cfg.save_dir, cfg.model_name)
|
||||
return load_checkpoint(build_model(cfg), path, device)
|
||||
|
||||
|
||||
def draw_prediction(label, image, threshold):
|
||||
"""Red: detected boxes, blue: cells with an object, green: empty cells."""
|
||||
boxes, cells_obj, cells_noobj = convert_prediction(label, image, threshold)
|
||||
if len(boxes) == 0:
|
||||
print("no boxes")
|
||||
image = draw_bounding_boxes(image, torch.tensor(cells_noobj), colors=(0, 255, 0))
|
||||
image = draw_bounding_boxes(image, torch.tensor(cells_obj), colors=(0, 0, 255))
|
||||
image = draw_bounding_boxes(image, torch.tensor(boxes), colors=(255, 0, 0))
|
||||
show_image(image)
|
||||
|
||||
|
||||
def train(cfg: Config):
|
||||
"""Train the detector; saves the model with the lowest validation loss."""
|
||||
device = get_device(cfg.device)
|
||||
train_loader = DataLoader(build_dataset(cfg, cfg.train_dir, cfg.augment_train),
|
||||
batch_size=cfg.batch_size, shuffle=True, num_workers=cfg.train_workers)
|
||||
val_loader = DataLoader(build_dataset(cfg, cfg.val_dir, cfg.augment_val),
|
||||
batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.val_workers)
|
||||
|
||||
model = build_model(cfg).to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=cfg.learning_rate, weight_decay=cfg.weight_decay)
|
||||
|
||||
fit(model, YoloLoss(), optimizer, train_loader, val_loader, device,
|
||||
epochs=cfg.epochs, save_dir=cfg.save_dir, model_name=cfg.model_name, save=cfg.save,
|
||||
monitor="loss")
|
||||
|
||||
|
||||
def view_data(cfg: Config):
|
||||
"""Show validation images with their ground-truth boxes, one at a time."""
|
||||
dataset = build_dataset(cfg, cfg.val_dir, cfg.augment_val)
|
||||
for images, labels in DataLoader(dataset, batch_size=1, shuffle=True):
|
||||
draw_prediction(labels[0], images[0], cfg.threshold)
|
||||
|
||||
|
||||
def sample(cfg: Config):
|
||||
"""Show the saved model's predictions on validation images."""
|
||||
device = get_device(cfg.device)
|
||||
model = load_trained(cfg, device)
|
||||
dataset = build_dataset(cfg, cfg.val_dir, cfg.augment_val)
|
||||
|
||||
for images, _ in DataLoader(dataset, batch_size=cfg.batch_size, shuffle=False):
|
||||
with torch.inference_mode():
|
||||
start = time.perf_counter()
|
||||
predictions = model(images.to(device)).cpu()
|
||||
print(f"Inference: {time.perf_counter() - start:.6f} seconds")
|
||||
for image, prediction in zip(images, predictions):
|
||||
draw_prediction(prediction, image, cfg.threshold)
|
||||
|
||||
|
||||
def benchmark_workers(cfg: Config):
|
||||
"""Time DataLoader num_workers values to pick the fastest."""
|
||||
device = get_device(cfg.device)
|
||||
dataset = build_dataset(cfg, cfg.train_dir, cfg.augment_train)
|
||||
benchmark_num_workers(dataset, build_model(cfg).to(device), device, batch_size=cfg.batch_size)
|
||||
|
||||
|
||||
def webcam(cfg: Config):
|
||||
"""Live detection on the webcam (--camera). Press q to quit."""
|
||||
import cv2
|
||||
|
||||
device = get_device(cfg.device)
|
||||
model = load_trained(cfg, device)
|
||||
|
||||
cap = cv2.VideoCapture(cfg.camera)
|
||||
if not cap.isOpened():
|
||||
raise RuntimeError("Could not open webcam")
|
||||
|
||||
def draw(frame, boxes, scale_w, scale_h, color, thickness):
|
||||
for b in boxes:
|
||||
cv2.rectangle(frame, (int(b[0] * scale_w), int(b[1] * scale_h)),
|
||||
(int(b[2] * scale_w), int(b[3] * scale_h)), color, thickness)
|
||||
|
||||
try:
|
||||
while True:
|
||||
ok, frame = cap.read()
|
||||
if not ok:
|
||||
break
|
||||
|
||||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
image = torch.from_numpy(rgb).permute(2, 0, 1).float() / 255.0
|
||||
_, H, W = image.shape
|
||||
scale_w, scale_h = W / cfg.image_size, H / cfg.image_size
|
||||
|
||||
image = F.interpolate(image.unsqueeze(0), size=(cfg.image_size, cfg.image_size),
|
||||
mode="bilinear", align_corners=True)
|
||||
with torch.inference_mode():
|
||||
prediction = model(image.to(device)).cpu()
|
||||
|
||||
boxes, cells_obj, cells_noobj = convert_prediction(prediction.squeeze(0), image.squeeze(0),
|
||||
threshold=cfg.threshold)
|
||||
draw(frame, cells_noobj, scale_w, scale_h, (0, 255, 0), 1)
|
||||
draw(frame, cells_obj, scale_w, scale_h, (255, 0, 0), 1)
|
||||
draw(frame, boxes, scale_w, scale_h, (0, 0, 255), 4)
|
||||
|
||||
cv2.imshow("Webcam", frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
finally:
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
TASKS = {
|
||||
"train": train,
|
||||
"view-data": view_data,
|
||||
"sample": sample,
|
||||
"benchmark-workers": benchmark_workers,
|
||||
"webcam": webcam,
|
||||
}
|
||||
@@ -1,160 +0,0 @@
|
||||
import os
|
||||
|
||||
import cv2
|
||||
from tqdm import tqdm
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import DataLoader, random_split
|
||||
from torchvision import datasets
|
||||
from torchvision.utils import draw_bounding_boxes
|
||||
|
||||
from src.util import TransformedSubset, test_workers_speed, visualizeImage
|
||||
from src.production.yolo_model_production import convert_prediction
|
||||
from src.yolo.yolo_dataset import YoloDataset
|
||||
from src.yolo.yolo_loss import YoloLoss
|
||||
from src.yolo.yolo_model import Yolo_model, sample, train
|
||||
|
||||
def view_data(dataset):
|
||||
dataloader = DataLoader(dataset=dataset, batch_size=1, shuffle=True)
|
||||
for images, labels in iter(dataloader):
|
||||
for batch in range(images.shape[0]):
|
||||
image = images[batch]
|
||||
label = labels[batch]
|
||||
visualize_boxes(label, image)
|
||||
|
||||
def sample_data(dataloader, model, device, SAVE_PATH):
|
||||
for images, labels in iter(dataloader):
|
||||
predictions = sample(model, images, device, SAVE_PATH)
|
||||
for batch in range(predictions.shape[0]):
|
||||
image = images[batch]
|
||||
prediction = predictions[batch]
|
||||
visualize_boxes(prediction, image)
|
||||
|
||||
def visualize_boxes(label, image, threshold=0.95):
|
||||
boxes_to_draw, grids_to_draw_obj, grids_to_draw_noobj = convert_prediction(label, image, threshold)
|
||||
if len(boxes_to_draw) == 0:
|
||||
print("no labels")
|
||||
boxes_to_draw = torch.tensor(boxes_to_draw)
|
||||
grids_to_draw_noobj = torch.tensor(grids_to_draw_noobj)
|
||||
grids_to_draw_obj = torch.tensor(grids_to_draw_obj)
|
||||
image = draw_bounding_boxes(image, grids_to_draw_noobj, colors=(0, 255, 0))
|
||||
image = draw_bounding_boxes(image, grids_to_draw_obj, colors=(0, 0, 255))
|
||||
image = draw_bounding_boxes(image, boxes_to_draw, colors=(255, 0, 0))
|
||||
visualizeImage(image)
|
||||
|
||||
def use_webcam(grid, img_size):
|
||||
# 0 = default webcam
|
||||
cap = cv2.VideoCapture(0)
|
||||
|
||||
if not cap.isOpened():
|
||||
raise Exception("Could not open webcam")
|
||||
|
||||
model = Yolo_model(c_in=3, boxes=1, grid=grid, labels=1)
|
||||
state_dict = torch.load(os.path.join("saved_models", "face_detection_yolo", "face_detection_yolo"), weights_only=False)
|
||||
model.load_state_dict(state_dict)
|
||||
model.eval()
|
||||
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
|
||||
if not ret:
|
||||
break
|
||||
|
||||
rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
|
||||
image = torch.from_numpy(rgb).permute(2, 0, 1).float()
|
||||
C, H, W = image.shape
|
||||
scale_w = W / img_size
|
||||
scale_h = H / img_size
|
||||
|
||||
image = image / 255.0
|
||||
|
||||
image = image.unsqueeze(0)
|
||||
image = F.interpolate(image, size=(img_size, img_size), mode="bilinear", align_corners=True) #different modes?
|
||||
|
||||
prediction = model(image)
|
||||
|
||||
bboxes, grid_ob, grid_noob = convert_prediction(prediction.squeeze(0), image.squeeze(0), threshold=0.9)
|
||||
|
||||
for bbox in grid_noob:
|
||||
xmin = int(bbox[0] * scale_w)
|
||||
ymin = int(bbox[1] * scale_h)
|
||||
xmax = int(bbox[2] * scale_w)
|
||||
ymax = int(bbox[3] * scale_h)
|
||||
cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 1)
|
||||
|
||||
for bbox in grid_ob:
|
||||
xmin = int(bbox[0] * scale_w)
|
||||
ymin = int(bbox[1] * scale_h)
|
||||
xmax = int(bbox[2] * scale_w)
|
||||
ymax = int(bbox[3] * scale_h)
|
||||
cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 0, 0), 1)
|
||||
|
||||
for bbox in bboxes:
|
||||
xmin = int(bbox[0] * scale_w)
|
||||
ymin = int(bbox[1] * scale_h)
|
||||
xmax = int(bbox[2] * scale_w)
|
||||
ymax = int(bbox[3] * scale_h)
|
||||
cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 0, 255), 4)
|
||||
|
||||
cv2.imshow("Webcam", frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cap.release()
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
|
||||
def train_yolo():
|
||||
SAVE_PATH = "./saved_models"
|
||||
IMAGE_SIZE = 64
|
||||
GRID = 6
|
||||
BATCH_SIZE = 64
|
||||
|
||||
dataset = YoloDataset(
|
||||
image_dir="data/faces_scenery/train",
|
||||
annotation_path="data/faces_scenery/train/_annotations.coco.json",
|
||||
img_size=IMAGE_SIZE,
|
||||
transform=True,
|
||||
grid=GRID
|
||||
)
|
||||
|
||||
dataset_valid = YoloDataset(
|
||||
image_dir="data/faces_scenery/test",
|
||||
annotation_path="data/faces_scenery/test/_annotations.coco.json",
|
||||
img_size=IMAGE_SIZE,
|
||||
transform=True,
|
||||
grid=GRID
|
||||
)
|
||||
|
||||
|
||||
train_loader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2)
|
||||
val_loader = DataLoader(dataset_valid, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)
|
||||
|
||||
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
|
||||
print("Using device", device)
|
||||
|
||||
model = Yolo_model(c_in=3, boxes=1, grid=GRID, labels=1)
|
||||
model.to(device)
|
||||
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
||||
loss_module = YoloLoss()
|
||||
|
||||
|
||||
view_data(dataset_valid)
|
||||
exit()
|
||||
|
||||
|
||||
#sample_data(dataloader=val_loader, model=model, device=device, SAVE_PATH=SAVE_PATH)
|
||||
#exit()
|
||||
|
||||
|
||||
#test_workers_speed(dataset, model)
|
||||
#exit()
|
||||
|
||||
|
||||
#use_webcam(GRID, IMAGE_SIZE)
|
||||
#exit()
|
||||
|
||||
|
||||
train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader,
|
||||
optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, model_name="face_detection_yolo")
|
||||
@@ -1,162 +0,0 @@
|
||||
import os
|
||||
import time
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
class Yolo_Conv_Block(nn.Module):
|
||||
def __init__(self, c_in, c_hidden, c_out, kernel_size):
|
||||
super().__init__()
|
||||
|
||||
self.model = nn.Sequential(
|
||||
nn.Conv2d(in_channels=c_in, out_channels=c_hidden, kernel_size=kernel_size, padding=kernel_size//2),
|
||||
nn.BatchNorm2d(c_hidden),
|
||||
nn.LeakyReLU(),
|
||||
nn.Conv2d(in_channels=c_hidden, out_channels=c_out, kernel_size=1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.model(x)
|
||||
|
||||
class SkipBlock(nn.Module):
|
||||
def __init__(self, c_in, c_out, kernel_size=3):
|
||||
super().__init__()
|
||||
self.conv = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, kernel_size, padding=kernel_size//2),
|
||||
nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
|
||||
nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
|
||||
nn.Conv2d(c_out, c_out, kernel_size, padding=kernel_size//2),
|
||||
nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
|
||||
nn.LeakyReLU(inplace=True)
|
||||
)
|
||||
self.conv_skip = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_out, 1),
|
||||
nn.GroupNorm(num_groups=c_out//8, num_channels=c_out),
|
||||
nn.LeakyReLU(inplace=True)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return(F.dropout(F.leaky_relu(self.conv_skip(x) + self.conv(x), inplace=True), p=0.3))
|
||||
|
||||
|
||||
class Yolo_model(nn.Module):
|
||||
def __init__(self, c_in, boxes, grid, labels, c_hidden=16):
|
||||
super().__init__()
|
||||
|
||||
self.model = nn.Sequential(
|
||||
nn.Conv2d(c_in, c_hidden, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(num_groups=c_hidden//8, num_channels=c_hidden),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden),
|
||||
|
||||
SkipBlock(c_in=c_hidden, c_out=c_hidden*2),
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*2),
|
||||
|
||||
SkipBlock(c_in=c_hidden*2, c_out=c_hidden*4),
|
||||
SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
|
||||
SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
|
||||
SkipBlock(c_in=c_hidden*4, c_out=c_hidden*4),
|
||||
|
||||
nn.Conv2d(c_hidden*4, c_hidden*8, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(num_groups=c_hidden//2, num_channels=c_hidden*8),
|
||||
nn.LeakyReLU(inplace=True),
|
||||
nn.Dropout(0.3),
|
||||
|
||||
nn.AdaptiveAvgPool2d((grid, grid)),
|
||||
nn.Conv2d(c_hidden*8, boxes*5 + labels, kernel_size=1)
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.model(x).permute(0, 2, 3, 1)
|
||||
return F.sigmoid(x)
|
||||
|
||||
|
||||
def train(model, loss_module, train_loader, val_loader, optimizer, SAVE_PATH, model_name, saving=True):
|
||||
best_val = torch.finfo(torch.float32).max
|
||||
device = next(model.parameters()).device
|
||||
|
||||
for epoch in range(500):
|
||||
############
|
||||
# Training #
|
||||
############
|
||||
model.train()
|
||||
|
||||
count, lossCount = 0, 0.
|
||||
for images, labels in tqdm(train_loader, desc=f"Train", leave=False):
|
||||
images = images.to(device)
|
||||
labels = labels.to(device)
|
||||
|
||||
prediction = model(images)
|
||||
|
||||
loss = loss_module(prediction, labels)
|
||||
lossCount += loss.item()
|
||||
|
||||
optimizer.zero_grad()
|
||||
|
||||
loss.backward()
|
||||
|
||||
optimizer.step()
|
||||
|
||||
count += images.size(0)
|
||||
|
||||
train_loss = lossCount / count
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
##############
|
||||
# Validation #
|
||||
##############
|
||||
model.eval()
|
||||
|
||||
count, lossCount = 0, 0.
|
||||
for images, labels in tqdm(val_loader, desc=f"Test", leave=False):
|
||||
with torch.no_grad():
|
||||
images = images.to(device)
|
||||
labels = labels.to(device)
|
||||
|
||||
prediction = model(images)
|
||||
|
||||
loss = loss_module(prediction, labels)
|
||||
lossCount += loss.item()
|
||||
|
||||
count += images.size(0)
|
||||
|
||||
val_loss = lossCount / count
|
||||
|
||||
if(saving and best_val > val_loss):
|
||||
best_val = val_loss
|
||||
save_dir = os.path.join(SAVE_PATH, model_name)
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
save_path = os.path.join(save_dir, model_name)
|
||||
torch.save(model.state_dict(), save_path)
|
||||
|
||||
print(f"epoch: {epoch+1} | train loss: {int(train_loss * 100000) / 100} | val loss: {int(val_loss * 100000) / 100}")
|
||||
torch.cuda.empty_cache()
|
||||
return best_val
|
||||
|
||||
|
||||
def sample(model, img, device, SAVE_PATH, model_name="face_detection_yolo", folder="face_detection_yolo"):
|
||||
with torch.no_grad():
|
||||
full_path = os.path.join(SAVE_PATH, folder, model_name)
|
||||
state_dict = torch.load(full_path, weights_only=False)
|
||||
model.load_state_dict(state_dict)
|
||||
model.eval()
|
||||
img = img.to(device)
|
||||
start = time.perf_counter()
|
||||
pred = model(img)
|
||||
end = time.perf_counter()
|
||||
print(f"Elapsed: {end - start:.6f} seconds")
|
||||
return pred
|
||||
Reference in New Issue
Block a user