From 7ca0bfa03c04f4a4baba6f7674403ede0b651375 Mon Sep 17 00:00:00 2001 From: yayoimizuha Date: Thu, 4 May 2023 20:14:50 +0900 Subject: [PATCH] update --- .drone.yml | 2 +- resnet_finetune_vggface.py | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.drone.yml b/.drone.yml index 64a20c4..b5ea4e4 100644 --- a/.drone.yml +++ b/.drone.yml @@ -22,4 +22,4 @@ steps: - mkdir -p data - $mount_command - ls data/ - - python resnet_finetune_vggface.py \ No newline at end of file + - CI=True python resnet_finetune_vggface.py \ No newline at end of file diff --git a/resnet_finetune_vggface.py b/resnet_finetune_vggface.py index 3b95750..a873b1a 100644 --- a/resnet_finetune_vggface.py +++ b/resnet_finetune_vggface.py @@ -1,4 +1,4 @@ -from os import makedirs +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, \ @@ -24,7 +24,9 @@ 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([ @@ -155,7 +157,7 @@ for epoch in range(epochs): model_gpu.train() makedirs(join(save_dir, 'pallets', str(epoch)), exist_ok=True) - for count, (images, labels) in enumerate(tqdm(dataloader['train'])): + for count, (images, labels) in enumerate(tqdm(dataloader['train'], disable=CI)): if count == 1: image_pallets = plot_dataset(dataloader=(images, labels), col_len=6, label_text=image_folder['train'].classes)