update

.drone.yml

@@ -20,4 +20,4 @@ steps:
       - mkdir -p data
       - $mount_command
       - ls data/
-      - python resnet_finetune.py
+      - python resnet_finetune_vggface.py

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)

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,26 +72,31 @@ 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)
+    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:
+            if res.__len__() == 1:
+                if res[0]['score'] == -1:
                     continue
 
             faces = []
             for face in res:
-            faces.append((face.kps, face.det_score, face.bbox))
+                faces.append((face['landmarks'], face['score'], face['bbox']))
             q2.put((faces, path))
 
 
@@ -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())

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

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

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

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

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

settings.py

@@ -64,5 +64,5 @@ request_header = {
 class FaceCropProcesses:
     load = 1
     pre_process = 10
-    predict = 3
+    predict = 2
     post_process = 4

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)):
+
+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(datadir(), 'sample_set', name, file),
-                     dst=join(datadir(), 'dataset', 'train', name, file))
+                # 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(datadir(), 'sample_set', name, file),
-                     dst=join(datadir(), 'dataset', 'val', name, file))
-
+                # 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))
+