tomokazu
/
helloproject-ai
mirror of https://github.com/yayoimizuha/helloproject-ai.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
helloproject-ai/convnext_finetune.py

232 lines
8.5 KiB
Python
Raw Blame History

from os import makedirs, environ
from torchinfo import summary
from torchvision.models import convnext_large, ConvNeXt_Large_Weights, convnext_base, ConvNeXt_Base_Weights
from torch.nn import Linear, Dropout3d, Sequential, Dropout
from torchvision.transforms import Compose, RandomResizedCrop, RandomRotation, ToTensor, \
RandomHorizontalFlip, \
Resize, CenterCrop, RandomAffine, GaussianBlur, RandomAutocontrast, InterpolationMode, AugMix, RandomErasing
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from numpy import arange, ndarray, ceil, full, uint8
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 distutils.util import strtobool
CI = bool(strtobool(environ['CI']))
device = 'cuda' if is_available() else 'cpu'
transform = {
'train': Compose([
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(),
RandomErasing(),
RandomRotation(degrees=15),
RandomResizedCrop(size=232, scale=(0.7, 1.0), ratio=(1.0, 1.0), antialias=True,
interpolation=InterpolationMode.BILINEAR),
# ConvNeXt_Large_Weights.IMAGENET1K_V1.transforms()
]),
'val': Compose([
# ConvNeXt_Large_Weights.IMAGENET1K_V1.transforms(),
RandomAffine(scale=(0.8, 0.8), degrees=(0, 0)),
Resize(232, antialias=True, interpolation=InterpolationMode.BILINEAR),
ToTensor()
])
}
image_folder = {
'train': ImageFolder(root=join(datadir(), 'dataset', 'train'), transform=transform['train']),
'val': ImageFolder(root=join(datadir(), 'dataset', '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 = convnext_base(weights=ConvNeXt_Base_Weights.IMAGENET1K_V1)
# model = resnet50(weights=None)
print()
tune = False
for name, layer in model.named_parameters():
if 'features.6' in name:
tune = True
layer.requires_grad = tune
# print(model)
model.classifier[2] = Linear(in_features=1024, out_features=image_folder['train'].classes.__len__(), bias=True)
# model.classifier = Sequential(Dropout(p=.5), model.classifier)
model.classifier.insert(0, Dropout3d(p=.5))
summary(model=model, input_size=(1, 3, 518, 518), 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.features[6].parameters(), 'lr': 1e-5},
{'params': model_gpu.features[7].parameters(), 'lr': 1e-4},
{'params': model_gpu.classifier.parameters(), 'lr': 1e-4},
])
# scheduler = lr_scheduler.StepLR(optimizer=optimizer, step_size=5, gamma=0.5)
epochs = 100
train_loss_list = list()
train_acc_list = list()
val_loss_list = list()
val_acc_list = list()
save_dir = join(datadir(), 'artifact', 'convnext-base_' + 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=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'))
Reference in New Issue View Git Blame Copy Permalink