diff --git a/.drone.yml b/.drone.yml index 52eef53..b1e1c2f 100644 --- a/.drone.yml +++ b/.drone.yml @@ -20,4 +20,4 @@ steps: - mkdir -p data - $mount_command - ls data/ - - python resnet_finetune.py \ No newline at end of file + - python resnet_finetune_vggface.py \ No newline at end of file diff --git a/deepface_test.py b/deepface_test.py new file mode 100644 index 0000000..a981a62 --- /dev/null +++ b/deepface_test.py @@ -0,0 +1,32 @@ +from PIL import Image, ImageDraw + +from matplotlib import pyplot +from retinaface.pre_trained_models import get_model +from torch import save +from retinaface.predict_single import Model +from retinaface.network import RetinaFace +from numpy import array + +im_path = "道重さゆみ=sayumimichishige-blog=12719215891-10.jpg" +model: Model = get_model(model_name='resnet50_2020-07-20', max_size=1024, device='cuda') +image = array(Image.open(im_path)) +faces = model.predict_jsons(image, confidence_threshold=0.4, nms_threshold=0.4) + +print(faces) +# print(dfs) +image = Image.open(im_path) +draw = ImageDraw.Draw(image) +if faces.__len__() == 1: + if faces[0]['score'] == -1: + exit() +for face in faces: + print(face['bbox']) + x, y, w, h = face['bbox'] + print(x, y, w, h) + draw.rectangle((x, y, w, h)) + for order, landmark in enumerate(face['landmarks'], start=1): + draw.point(tuple(landmark)) + draw.text(tuple(landmark), text=str(order)) +# +image.save("dest.png") +# RetinaFace.detect_faces(im_path) diff --git a/face_crop.py b/face_crop.py index ea90eb8..febec81 100755 --- a/face_crop.py +++ b/face_crop.py @@ -1,22 +1,21 @@ -import torch.cuda -from numpy import ndarray -import numpy as np +import torch.backends.cudnn +from torch import no_grad, backends import settings from os import makedirs, listdir, stat, utime, devnull from os.path import join, exists from tqdm import tqdm -from PIL import Image, ImageDraw -from numpy import array, arctan2, pi, zeros, uint8, float32 +from PIL import Image +from numpy import array, arctan2, pi from aiofiles import open as a_open from asyncio import gather, run from multiprocessing import Queue, Process, get_start_method, set_start_method -from time import time, sleep +from time import sleep from io import BytesIO -from math import ceil, sqrt -from torch import from_numpy, cuda, Tensor, inference_mode, nn -import atexit -from insightface.app import FaceAnalysis -from contextlib import redirect_stdout +from math import sqrt +# from insightface.app import FaceAnalysis +from retinaface.pre_trained_models import get_model +from retinaface.predict_single import Model +from contextlib import redirect_stdout, redirect_stderr from collections import OrderedDict face_dir = join(settings.datadir(), 'face_cropped') @@ -37,7 +36,7 @@ def truncate(landmark: list[tuple[float]]) -> tuple[tuple[int, int], float]: def load_image(basedir: str, queue: Queue, progress: tuple[Queue]) -> None: def list_up(): - for name in ["田中れいな", "田中れいな"]: # listdir(basedir): + for name in listdir(basedir): # ["江端妃咲"]: for image_file in listdir(join(basedir, name)): yield name, image_file @@ -52,7 +51,7 @@ def load_image(basedir: str, queue: Queue, progress: tuple[Queue]) -> None: bar = tqdm(total=file_list.__len__()) for i in range(0, file_list.__len__(), 20): - while queue.qsize() > 150: + while queue.qsize() > 50: sleep(1e-3) chunk = file_list[i:i + 20] @@ -73,27 +72,32 @@ def pre_process(q1: Queue, q2: Queue): # sleep(1e-4) image, path = q1.get() # print(type(image)) - img_arr = array(image)[:, :, ::-1] + img_arr = array(image) # [:, :, ::-1] q2.put((img_arr, path)) def predict(q1: Queue, q2: Queue, gpu: int): sleep(gpu * 4) - with redirect_stdout(open(devnull, mode='w')): - face_analysis = FaceAnalysis(providers=['CUDAExecutionProvider'], allowed_modules=['detection']) - face_analysis.prepare(ctx_id=gpu) - while True: - image, path = q1.get() + torch.backends.cudnn.benchmark = True + with redirect_stderr(open(devnull, mode='w')): + # face_analysis = FaceAnalysis(providers=['CUDAExecutionProvider'], allowed_modules=['detection']) + # face_analysis.prepare(ctx_id=gpu) + model: Model = get_model(model_name='resnet50_2020-07-20', max_size=512, device='cuda') + model.eval() + with no_grad(): + while True: + image, path = q1.get() - res = face_analysis.get(image) + res = model.predict_jsons(image=image, confidence_threshold=0.4, nms_threshold=0.4) - if not res: - continue + if res.__len__() == 1: + if res[0]['score'] == -1: + continue - faces = [] - for face in res: - faces.append((face.kps, face.det_score, face.bbox)) - q2.put((faces, path)) + faces = [] + for face in res: + faces.append((face['landmarks'], face['score'], face['bbox'])) + q2.put((faces, path)) def post_process(queue: Queue): @@ -103,15 +107,15 @@ def post_process(queue: Queue): image = Image.open(join(blog_images, *path)) name, file = path - width, height = image.size + # width, height = image.size # if width * height > 400_0000: # image = image.resize(size=(width // 2, height // 2)) for order, face in enumerate(res): - kps, score, bbox = face + landmarks, score, bbox = face face_width = bbox[2] - bbox[0] face_height = bbox[3] - bbox[1] - trans = truncate(kps) + trans = truncate(landmarks) rotated = image.rotate(angle=trans[1] * 360 / (2 * pi), center=trans[0]) image_size = max(face_width, face_height) * sqrt(2) // 2 if image_size < 100: @@ -145,6 +149,7 @@ if __name__ == '__main__': [p.start() for p in PreProcesses] [p.start() for p in Predict_Process] [p.start() for p in PostProcesses] + sleep(20) while True: sleep(5) # print(Load_Q.qsize(), PreProcess_Q.qsize(), Predict_Q.qsize(), PostProcess_Q.qsize()) diff --git a/inceptionnet_finetune.py b/inceptionnet_finetune.py index 4f5065e..4161a66 100644 --- a/inceptionnet_finetune.py +++ b/inceptionnet_finetune.py @@ -1,6 +1,7 @@ from os import makedirs from torchvision.models import Inception_V3_Weights, inception_v3 from torchvision.models import swin_v2_b, Swin_V2_B_Weights +from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights from torch.nn import Linear from torchvision.transforms import Compose, RandomResizedCrop, RandomRotation, ToTensor, \ RandomHorizontalFlip, \ @@ -102,7 +103,9 @@ for i in range(epochs): images = images.to(device) labels = labels.to(device) - outputs, _ = model_gpu(images) + o = model_gpu(images) + print(o) + outputs, _ = o loss = criterion(outputs, labels) train_loss += loss.item() diff --git a/mask_correct.py b/mask_correct.py new file mode 100644 index 0000000..2577c3e --- /dev/null +++ b/mask_correct.py @@ -0,0 +1,40 @@ +from torch import zeros, load, no_grad, stack, float32, nn +from torch.utils.data import DataLoader +from torchvision.datasets import ImageFolder +from torchvision.transforms import Compose, functional, ToTensor, Resize, ConvertImageDtype +from torch2trt import torch2trt +from PIL import Image +from numpy import array +from os import listdir, makedirs +from settings import datadir +from os.path import join, isdir, isfile, basename +from tqdm import tqdm +from more_itertools import chunked +from shutil import copyfile + +transform = Compose([ + ToTensor(), + # ConvertImageDtype(float32), + Resize(224, antialias=True) +]) + +image_folder = ImageFolder(root=join(datadir(), 'face_cropped'), transform=transform) +print(image_folder.__len__()) +dataloader = DataLoader(image_folder, batch_size=300, shuffle=False, num_workers=14, pin_memory=True) + +dest_dir = join(datadir(), 'masked_or_not', 'infer') +makedirs(dest_dir, exist_ok=True) +mask_status = ('no', 'yes') +[makedirs(join(dest_dir, p)) for p in mask_status] +[makedirs(join(dest_dir, p, name)) for p in mask_status for name in image_folder.classes] + +model = load(join(datadir(), 'artifact', 'masked_or_not_2023-04-30 19:38:14.398157/model.pth')).cuda() +model.eval() +with no_grad(): + for (images, labels), fileinfo in zip(tqdm(dataloader), chunked(image_folder.imgs, n=300)): + # print(fileinfo) + res = model(images.cuda()) + for is_mask, (filename, person) in zip(res.max(1).indices.tolist(), fileinfo): + # print(is_mask,image_folder.classes[person],filename) + copyfile(filename, join(dest_dir, mask_status[is_mask], image_folder.classes[person], basename(filename))) + # break diff --git a/masked_or_not.py b/masked_or_not.py new file mode 100644 index 0000000..5ab9039 --- /dev/null +++ b/masked_or_not.py @@ -0,0 +1,232 @@ +from os import makedirs +import numpy +import torch.backends.cudnn +from torchvision.models import mobilenet_v3_small, MobileNet_V3_Small_Weights +from torch.nn import Linear +from torchvision.transforms import Compose, RandomResizedCrop, RandomRotation, ToTensor, \ + RandomHorizontalFlip, \ + Resize, RandomAffine, RandomAdjustSharpness, RandomAutocontrast, RandomEqualize, GaussianBlur, Normalize +from numpy import arange, ceil, full, float32, uint8, amax, amin +from torchsummary import summary +from torch.nn import CrossEntropyLoss +from torch.cuda.amp.grad_scaler import GradScaler +from torch.cuda.amp import autocast + +from torch.optim import SGD, Adam, lr_scheduler +from torchvision.datasets import ImageFolder +from torch.utils.data import DataLoader +from tqdm import tqdm +from settings import datadir +from os.path import join +from torch.cuda import is_available +from torch import no_grad, save +from datetime import datetime +from PIL import Image, ImageDraw, ImageFont +import matplotlib + +matplotlib.use('Agg') +device = 'cuda' if is_available() else 'cpu' +torch.backends.cudnn.deterministic = True + +transform = { + 'train': Compose([ + # CenterCrop(200), + RandomHorizontalFlip(p=0.1), + RandomAdjustSharpness(sharpness_factor=2, p=0.2), + GaussianBlur(kernel_size=3), + RandomAutocontrast(), + RandomEqualize(p=0.5), + RandomRotation(degrees=15), + ToTensor(), + # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + Resize(size=224, antialias=True) + ]), + 'val': Compose([ + # CenterCrop(200), + ToTensor(), + # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + Resize(size=224, antialias=True) + ]) +} + +image_folder = { + 'train': ImageFolder(root=join(datadir(), 'masked_or_not', 'train'), transform=transform['train']), + 'val': ImageFolder(root=join(datadir(), 'masked_or_not', 'val'), transform=transform['val']) +} + +dataloader = { + 'train': DataLoader(image_folder['train'], batch_size=16, shuffle=True, num_workers=8), + 'val': DataLoader(image_folder['val'], batch_size=16, shuffle=False, num_workers=8) +} + + +def plot_dataset(dataloader: DataLoader | tuple, col_len: int = 8, + label_text: str | None = None) -> Image.Image: + if isinstance(dataloader, DataLoader): + images, labels = iter(dataloader).__next__() + else: + images, labels = dataloader + + images: torch.Tensor = images + labels: torch.Tensor = labels + images: numpy.ndarray = images.numpy() + + if label_text is None: + labels: list[str] = [str(i) for i in labels.tolist()] + else: + labels: list[str] = [label_text[i] for i in labels.tolist()] + + batch_size, _, width, height = images.shape + # print(batch_size, width, height) + # print(images.dtype) + rows = ceil(batch_size / col_len) + # print(amax(images), amin(images)) + space_y, space_x, font_size = 50, 30, 20 + shape_y, shape_x = images.shape[-2:] + base_img = full(shape=((height + space_y) * int(rows), width * col_len + space_x * (col_len - 1), 3), dtype=uint8, + fill_value=255) + for order, image in enumerate(images): + order_y, order_x = order // col_len, order % col_len + image = (image.transpose([1, 2, 0]) * 255).astype(uint8) + # print(order_y, order_x) + # print(order_y * (shape_y + 30) + 30, (order_y + 1) * (shape_y + 30), + # order_x * (shape_x + 20), (order_x + 1) * (shape_x + 20) - 20) + base_img[order_y * (shape_y + space_y) + space_y:(order_y + 1) * (shape_y + space_y), + order_x * (shape_x + space_x):(order_x + 1) * (shape_x + space_x) - space_x, :] = image + pil_image = Image.fromarray(base_img) + font = ImageFont.truetype(font='NotoSansCJK-Medium.ttc', size=24) + draw = ImageDraw.Draw(pil_image) + pad = 5 + for order, label in enumerate(labels): + order_y, order_x = order // col_len, order % col_len + draw.text(((shape_x + space_x) * order_x + pad, (shape_y + space_y) * order_y + pad), label, 'black', font=font) + + return pil_image + # pyplot.imshow((images[0].transpose([1, 2, 0]) * 255).astype(uint8)) + + +model = mobilenet_v3_small(weights=MobileNet_V3_Small_Weights.IMAGENET1K_V1) + +# print(model) +summary(model=model, input_size=(3, 224, 224), device='cpu') + +for layer_name, layer in model.named_parameters(): + # print(layer_name) + layer.requires_grad = False + if 'classifier' in layer_name: + layer.requires_grad = True + if 'features.12' in layer_name: + layer.requires_grad = True + if 'features.11' in layer_name: + layer.requires_grad = True + if 'features.10' in layer_name: + layer.requires_grad = True + +model.classifier[3] = Linear(in_features=1024, out_features=image_folder['train'].classes.__len__(), bias=True) + +model_gpu = model.cuda() +criterion = CrossEntropyLoss().cuda() + +optimizer = Adam(params=[ + {'params': model_gpu.classifier.parameters(), 'lr': 1e-3}, + {'params': model_gpu.features[12].parameters(), 'lr': 1e-4}, + {'params': model_gpu.features[11].parameters(), 'lr': 1e-4}, + {'params': model_gpu.features[10].parameters(), 'lr': 1e-4}, +]) + +scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=5, gamma=.5) +epochs = 50 + +train_loss_list = list() +train_acc_list = list() +val_loss_list = list() +val_acc_list = list() + +save_dir = join(datadir(), 'artifact', 'masked_or_not_' + datetime.now().__str__()) +print(save_dir) +makedirs(save_dir, exist_ok=True) +makedirs(join(save_dir, 'pallets'), exist_ok=True) + +scaler = GradScaler() + +for epoch in range(epochs): + train_loss = .0 + train_acc = .0 + val_loss = .0 + val_acc = .0 + + model_gpu.train() + makedirs(join(save_dir, 'pallets', str(epoch)), exist_ok=True) + for count, (images, labels) in enumerate(tqdm(dataloader['train'])): + image_pallets = plot_dataset(dataloader=(images, labels), col_len=4, label_text=image_folder['train'].classes) + image_pallets.save(join(save_dir, 'pallets', str(epoch), str(count) + '.jpg')) + optimizer.zero_grad() + images = images.cuda() + labels = labels.cuda() + with autocast(): + outputs = model_gpu(images) + loss = criterion(outputs, labels) + train_loss += loss.item() + + # loss.backward() + # optimizer.step() + scaler.scale(loss).backward() + scaler.step(optimizer=optimizer) + scaler.update() + predicted = outputs.max(1)[1] + train_acc += (predicted == labels).sum() + + avg_train_loss = train_loss / dataloader['train'].dataset.__len__() + avg_train_acc = train_acc / dataloader['train'].dataset.__len__() + + model_gpu.eval() + with no_grad(): + for images, labels in dataloader['val']: + images = images.to(device) + labels = labels.to(device) + outputs = model_gpu(images) + loss = criterion(outputs, labels) + val_loss += loss.item() + predicted = outputs.max(1)[1] + val_acc += (predicted == labels).sum() + avg_val_loss = val_loss / dataloader['val'].dataset.__len__() + avg_val_acc = val_acc / dataloader['val'].dataset.__len__() + + print(f'Epoch [{(epoch + 1):02}/{epochs}], loss: {avg_train_loss:.5f}, ' + f'acc: {avg_train_acc:.5f}, val_loss: {avg_val_loss:.5f}, val_acc: {avg_val_acc:.5f}, ' + f'lr: {scheduler.get_last_lr()[0]:.2e}') + scheduler.step() + + train_loss_list.append(float(avg_train_loss)) + train_acc_list.append(float(avg_train_acc)) + val_loss_list.append(float(avg_val_loss)) + val_acc_list.append(float(avg_val_acc)) + + import matplotlib.pyplot as plt + + plt.figure(figsize=(8, 6)) + plt.plot(val_acc_list, label='val', lw=2, c='b') + plt.plot(train_acc_list, label='train', lw=2, c='k') + plt.title('learning rate') + plt.xticks(size=14) + plt.yticks(size=14) + plt.grid(lw=2) + plt.legend(fontsize=14, bbox_to_anchor=(1, 0)) + plt.xticks(arange(1, epochs, 2)) + plt.savefig(join(save_dir, 'learning_rate.png')) + plt.close() + + plt.figure(figsize=(8, 6)) + plt.plot(val_loss_list, label='val', lw=2, c='b') + plt.plot(train_loss_list, label='train', lw=2, c='k') + plt.title('loss') + plt.xticks(size=14) + plt.yticks(size=14) + plt.grid(lw=2) + plt.legend(fontsize=14, bbox_to_anchor=(1, 1)) + plt.xticks(arange(1, epochs, 2)) + plt.savefig(join(save_dir, 'loss.png')) + plt.close() + del plt + +save(model_gpu.cpu(), join(save_dir, 'model.pth')) diff --git a/resnet_finetune.py b/resnet_finetune.py index df29c91..5b3a74d 100644 --- a/resnet_finetune.py +++ b/resnet_finetune.py @@ -102,7 +102,7 @@ for i in range(epochs): images = images.to(device) labels = labels.to(device) - outputs = model_gpu(images) + outputs = model(images) loss = criterion(outputs, labels) train_loss += loss.item() diff --git a/resnet_finetune_vggface.py b/resnet_finetune_vggface.py new file mode 100644 index 0000000..4f4e0e9 --- /dev/null +++ b/resnet_finetune_vggface.py @@ -0,0 +1,223 @@ +from os import makedirs +from torchvision.models import ResNet50_Weights, resnet50 +from torch.nn import Linear +from torchvision.transforms import Compose, RandomResizedCrop, RandomRotation, ToTensor, \ + RandomHorizontalFlip, \ + Resize, CenterCrop, RandomAffine, GaussianBlur, RandomAutocontrast +import matplotlib +matplotlib.use('Agg') +import matplotlib.pyplot as plt +from numpy import arange, ndarray, ceil, full, uint8 +from torchsummary import summary +from torch.nn import CrossEntropyLoss +from torch.optim import SGD, Adam, lr_scheduler +from torchvision.datasets import ImageFolder +from torch.utils.data import DataLoader +from tqdm import tqdm +from PIL import Image, ImageDraw, ImageFont +from settings import datadir +from os.path import join +from torch.cuda import is_available +from torch import no_grad, save, Tensor +from datetime import datetime + +device = 'cuda' if is_available() else 'cpu' +transform = { + 'train': Compose([ + Resize(448), + CenterCrop(224), + RandomHorizontalFlip(p=0.1), + GaussianBlur(kernel_size=3), + RandomAutocontrast(), + # Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ToTensor(), + RandomRotation(degrees=15), + RandomResizedCrop(size=224, scale=(0.7, 1.0), ratio=(1.0, 1.0), antialias=True) + ]), + 'val': Compose([ + Resize(448), + CenterCrop(224), + ToTensor(), + RandomAffine(scale=(0.8, 0.8), degrees=(0, 0)), + # Resize(224, antialias=True) + ]) +} +image_folder = { + 'train': ImageFolder(root=join(datadir(), 'vggface2', 'train'), transform=transform['train']), + 'val': ImageFolder(root=join(datadir(), 'vggface2', 'val'), transform=transform['val']) +} + +dataloader = { + 'train': DataLoader(image_folder['train'], batch_size=64, shuffle=True, num_workers=3), + 'val': DataLoader(image_folder['val'], batch_size=64, shuffle=False, num_workers=3) +} + + +def plot_dataset(dataloader: DataLoader | tuple, col_len: int = 8, + label_text: str | None = None) -> Image.Image: + if isinstance(dataloader, DataLoader): + images, labels = iter(dataloader).__next__() + else: + images, labels = dataloader + + images: Tensor = images + labels: Tensor = labels + images: ndarray = images.numpy() + + if label_text is None: + labels: list[str] = [str(i) for i in labels.tolist()] + else: + labels: list[str] = [label_text[i] for i in labels.tolist()] + + batch_size, _, width, height = images.shape + # print(batch_size, width, height) + # print(images.dtype) + rows = ceil(batch_size / col_len) + # print(amax(images), amin(images)) + space_y, space_x, font_size = 50, 30, 20 + shape_y, shape_x = images.shape[-2:] + base_img = full(shape=((height + space_y) * int(rows), width * col_len + space_x * (col_len - 1), 3), dtype=uint8, + fill_value=255) + for order, image in enumerate(images): + order_y, order_x = order // col_len, order % col_len + image = (image.transpose([1, 2, 0]) * 255).astype(uint8) + # print(order_y, order_x) + # print(order_y * (shape_y + 30) + 30, (order_y + 1) * (shape_y + 30), + # order_x * (shape_x + 20), (order_x + 1) * (shape_x + 20) - 20) + base_img[order_y * (shape_y + space_y) + space_y:(order_y + 1) * (shape_y + space_y), + order_x * (shape_x + space_x):(order_x + 1) * (shape_x + space_x) - space_x, :] = image + pil_image = Image.fromarray(base_img) + font = ImageFont.truetype(font='NotoSansCJK-Medium.ttc', size=24) + draw = ImageDraw.Draw(pil_image) + pad = 5 + for order, label in enumerate(labels): + order_y, order_x = order // col_len, order % col_len + draw.text(((shape_x + space_x) * order_x + pad, (shape_y + space_y) * order_y + pad), label, 'black', font=font) + + return pil_image + # pyplot.imshow((images[0].transpose([1, 2, 0]) * 255).astype(uint8)) + + +model = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2) +# model = resnet50(weights=None) + +print() + +tune = False +for name, layer in model.named_parameters(): + if 'layer3' in name: + tune = True + layer.requires_grad = tune + +# print(model) + +model.fc = Linear(in_features=2048, out_features=image_folder['train'].classes.__len__(), bias=True) +# summary(model=model, input_size=(3, 224, 224), device='cpu') + +model_gpu = model.to(device=device) +criterion = CrossEntropyLoss() +# optimizer = Adam(model_gpu.parameters(), lr=1e-4) +optimizer = Adam(params=[ + # {'params': model_gpu.conv1.parameters(), 'lr': 1e-8}, + # {'params': model_gpu.bn1.parameters(), 'lr': 1e-8}, + # {'params': model_gpu.relu.parameters(), 'lr': 1e-8}, + # {'params': model_gpu.maxpool.parameters(), 'lr': 1e-8}, + {'params': model_gpu.layer1.parameters(), 'lr': 1e-8}, + {'params': model_gpu.layer2.parameters(), 'lr': 1e-8}, + {'params': model_gpu.layer3.parameters(), 'lr': 1e-5}, + {'params': model_gpu.layer4.parameters(), 'lr': 1e-4}, + {'params': model_gpu.fc.parameters(), 'lr': 1e-4}, + +]) +scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=5, gamma=0.5) +epochs = 50 + +train_loss_list = list() +train_acc_list = list() +val_loss_list = list() +val_acc_list = list() + +save_dir = join(datadir(), 'artifact', 'vggface2_' + datetime.now().__str__()) +print(save_dir) +makedirs(save_dir, exist_ok=True) +makedirs(join(save_dir, 'pallets'), exist_ok=True) + +for epoch in range(epochs): + train_loss = .0 + train_acc = .0 + val_loss = .0 + val_acc = .0 + + model_gpu.train() + makedirs(join(save_dir, 'pallets', str(epoch)), exist_ok=True) + + for count, (images, labels) in enumerate(tqdm(dataloader['train'])): + if count == 1: + image_pallets = plot_dataset(dataloader=(images, labels), col_len=8, label_text=image_folder['train'].classes) + image_pallets.save(join(save_dir, 'pallets', str(epoch), 'pallet.jpg')) + optimizer.zero_grad() + images = images.to(device) + labels = labels.to(device) + + outputs = model(images) + + loss = criterion(outputs, labels) + train_loss += loss.item() + + loss.backward() + optimizer.step() + + predicted = outputs.max(1)[1] + train_acc += (predicted == labels).sum() + + avg_train_loss = train_loss / dataloader['train'].dataset.__len__() + avg_train_acc = train_acc / dataloader['train'].dataset.__len__() + + model_gpu.eval() + with no_grad(): + for images, labels in dataloader['val']: + images = images.to(device) + labels = labels.to(device) + outputs = model_gpu(images) + loss = criterion(outputs, labels) + val_loss += loss.item() + predicted = outputs.max(1)[1] + val_acc += (predicted == labels).sum() + avg_val_loss = val_loss / dataloader['val'].dataset.__len__() + avg_val_acc = val_acc / dataloader['val'].dataset.__len__() + + print(f'Epoch [{(epoch + 1):02}/{epochs}], loss: {avg_train_loss:.5f}, ' + f'acc: {avg_train_acc:.5f}, val_loss: {avg_val_loss:.5f}, val_acc: {avg_val_acc:.5f}, ' + f'lr: {scheduler.get_last_lr()[0]:.2e}') + scheduler.step() + + train_loss_list.append(float(avg_train_loss)) + train_acc_list.append(float(avg_train_acc)) + val_loss_list.append(float(avg_val_loss)) + val_acc_list.append(float(avg_val_acc)) + + plt.figure(figsize=(8, 6)) + plt.plot(val_acc_list, label='val', lw=2, c='b') + plt.plot(train_acc_list, label='train', lw=2, c='k') + plt.title('learning rate') + plt.xticks(size=14) + plt.yticks(size=14) + plt.grid(lw=2) + plt.legend(fontsize=14) + plt.xticks(arange(0, epochs, 2)) + plt.savefig(join(save_dir, 'learning_rate.png')) + plt.close() + + plt.figure(figsize=(8, 6)) + plt.plot(val_loss_list, label='val', lw=2, c='b') + plt.plot(train_loss_list, label='train', lw=2, c='k') + plt.title('loss') + plt.xticks(size=14) + plt.yticks(size=14) + plt.grid(lw=2) + plt.legend(fontsize=14) + plt.xticks(arange(0, epochs, 2)) + plt.savefig(join(save_dir, 'loss.png')) + plt.close() + +save(model_gpu.cpu(), join(save_dir, 'model.pth')) diff --git a/settings.py b/settings.py index e40fddd..ed38587 100755 --- a/settings.py +++ b/settings.py @@ -64,5 +64,5 @@ request_header = { class FaceCropProcesses: load = 1 pre_process = 10 - predict = 3 + predict = 2 post_process = 4 diff --git a/split_train_val.py b/split_train_val.py index 94a73db..b3fd121 100644 --- a/split_train_val.py +++ b/split_train_val.py @@ -3,25 +3,52 @@ from os.path import join from settings import datadir from shutil import rmtree, copyfile from random import random +from tqdm import tqdm +from collections import OrderedDict +from asyncio import to_thread, gather, run +from aiofiles import open as a_open valid_rate = 0.1 +SRC_DIR = join(r'/mnt/share/dataset/vggface2/train') +DEST_DIR = join(datadir(), 'vggface2') -makedirs(join(datadir(), 'dataset'), exist_ok=True) -rmtree(join(datadir(), 'dataset', 'train'), ignore_errors=True) -rmtree(join(datadir(), 'dataset', 'val'), ignore_errors=True) -makedirs(join(datadir(), 'dataset', 'train'), exist_ok=True) -makedirs(join(datadir(), 'dataset', 'val'), exist_ok=True) +makedirs(DEST_DIR, exist_ok=True) +rmtree(join(DEST_DIR, 'train'), ignore_errors=True) +rmtree(join(DEST_DIR, 'val'), ignore_errors=True) +makedirs(join(DEST_DIR, 'train'), exist_ok=True) +makedirs(join(DEST_DIR, 'val'), exist_ok=True) -for name in listdir(join(datadir(), 'sample_set')): - print(name) - makedirs(join(datadir(), 'dataset', 'train', name)) - makedirs(join(datadir(), 'dataset', 'val', name)) - for file in listdir(join(datadir(), 'sample_set', name)): - if random() > valid_rate: - copyfile(src=join(datadir(), 'sample_set', name, file), - dst=join(datadir(), 'dataset', 'train', name, file)) - else: - copyfile(src=join(datadir(), 'sample_set', name, file), - dst=join(datadir(), 'dataset', 'val', name, file)) - # print(name, file) +async def waiting(arg): + return await gather(*arg) + + +async def async_copyfile(src_path: str, dst_path: str): + async with a_open(file=src_path, mode='rb') as f: + cont = await f.read() + async with a_open(file=dst_path, mode='wb') as f: + await f.write(cont) + + +with tqdm(listdir(SRC_DIR)) as pbar: + for name in pbar: + pbar.set_postfix(OrderedDict(name=name)) + # print(name) + makedirs(join(DEST_DIR, 'train', name)) + makedirs(join(DEST_DIR, 'val', name)) + coroutines = [] + for file in listdir(join(SRC_DIR, name)): + if random() > valid_rate: + # copyfile(src=join(SRC_DIR, name, file), + # dst=join(DEST_DIR, 'train', name, file)) + # co = to_thread(copyfile, join(SRC_DIR, name, file), join(DEST_DIR, 'train', name, file)) + co = async_copyfile(join(SRC_DIR, name, file), join(DEST_DIR, 'train', name, file)) + else: + # copyfile(src=join(SRC_DIR, name, file), + # dst=join(DEST_DIR, 'val', name, file)) + # co = to_thread(copyfile, join(SRC_DIR, name, file), join(DEST_DIR, 'val', name, file)) + co = async_copyfile(join(SRC_DIR, name, file), join(DEST_DIR, 'val', name, file)) + coroutines.append(co) + # print(name, file) + run(waiting(coroutines)) +