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