231 lines
8.3 KiB
Python
231 lines
8.3 KiB
Python
from os import makedirs, environ
|
|
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
|
|
from pytorch_metric_learning.losses import ArcFaceLoss
|
|
from pytorch_metric_learning.distances import CosineSimilarity
|
|
from pytorch_metric_learning.regularizers import RegularFaceRegularizer
|
|
from distutils.util import strtobool
|
|
|
|
CI = bool(strtobool(environ['CI']))
|
|
device = 'cuda' if is_available() else 'cpu'
|
|
transform = {
|
|
'train': Compose([
|
|
Resize(350),
|
|
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(350),
|
|
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=32, shuffle=True, num_workers=3),
|
|
'val': DataLoader(image_folder['val'], batch_size=32, 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=r'/usr/share/fonts/opentype/noto/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 'layer1' 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-6},
|
|
{'params': model_gpu.layer4.parameters(), 'lr': 1e-5},
|
|
{'params': model_gpu.fc.parameters(), 'lr': 1e-5},
|
|
|
|
])
|
|
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'], disable=CI, mininterval=20)):
|
|
if count == 1:
|
|
image_pallets = plot_dataset(dataloader=(images, labels), col_len=6,
|
|
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'))
|