github의 많은 코드들 중에서 star가 많은 별들의 코드를 scatch 하며 하나씩 파악해나가는 것은 매우 중요하다고 생각한다. 현재 트렌드의 코드 기법과 해당 개발자의 의도를 파악할 수 있기 때문이다.
이번에는 PFLD의 train 과정에 대해서 설명해보려고 한다. 또한, 저자는 loss 비교를 위해 tensorboard를 사용하였지만 나는 Weight & Bias를 사용하려고 한다. tensorboard보다 visualize하는데 편하고 좀더 관리하기에 용이하다고 생각한다.
해당 코드를 scratch하면서 기본적인 내용들도 같이 설명하려고 한다. 내용이 길어질 수 있겠으나 복습차원으로도 공부가 될 것이고, 그냥 넘어갔던 내용들이 복기되면서 나의 코딩 스킬이 늘어날 수 있을거라 생각하여 같이 작성해보려고 한다.
우선은 argparse에 대해 먼저 알아보려고 한다.
import argparse
우리는 흔이 이처럼 라이브러리를 불러와 사용하는데, window나 linux의 command, terminal에서 yolov4 yolov5, yolov8을 실행했을때 여러 명령어와 인자를 넣어서 사용하곤 했다. 그럼 이게 도대체 무엇인지에 대해서는 공부를 해봤는가?
그냥 argparse를 사용하면 command, terminal단에서 쉽게 조정이 가능한것만 알면 될까? 솔직히 그렇게만 알아도 무방하다고 생각한다. 하지만 앞서 말했듯이 복습, 복기 하는 과정을 걸쳐서 좀더 나은 코딩스킬을 쌓기 위해 기본 개념을 알아가보자.
argparse란 프로그램에 필요한 인자를 사용자 친화적인 명령행 인터페이스로 쉽게 작성하도록 돕는 라이브러리이다. 즉, command 창에서 프로그램 내의 인자를 조절하게끔 도와준다. 참고로 argparse는 python에 기본적으로 내장되어 있다.
해당 내용은 구글링 시 처음 나오는 내용인데, 좀 어렵게 말하면 위와 같은 정의이고 쉽게 말하면 command, terminal 창에서 인자를 집어넣어 코드를 실행할 수 있는 구조이다.
우리가 학습을 할때, dataset의 위치, batch_size, epoch, image_size 등 바꿔야할 파라미터들이 많다. 학습할 모든 데이터셋들이 특정 개발자가 오픈소스로 배포한 코드의 변수에 모두 같지 않으므로 다를 수밖에 없는 부분들은 따로 빼서 argparse로 사용자마다 다양한 인자를 받아 변수에 넣을 수 있도록 도와주는 라이브러리라고 생각하면 된다.
import argparse
parser = argparse.ArgumentParser(description="PFLD")
parser.add_argument('--base_lr', default=0.0001, type=int)
parser.add_argument('--weight-decay', '--wd', default=1e-6, type=float)
# -- lr
parser.add_argument("--lr_patience", default=40, type=int)
# -- epoch
parser.add_argument('--start_epoch', default=1, type=int)
parser.add_argument('--end_epoch', default=500, type=int)
argparse.ArgumentParser() 클래스는 입력받은 인자를 파싱하고 처리하는데 사용되는데, description으로 해당 Parser가 어떤 작업을 하는지 간략하게 작성할 수 있는 기능이 있다. description이 없어도 상관없으니 가독성 측면에선 간략하게 작성해주는 것이 좋다.
logging은 오류의 발생 위치 및 종류를 확인하는 라이브러리이다. 과정에 대해서 상세히 설명하는 블로그는 https://sosoeasy.tistory.com/414 여기서 참고하면 좋다.
logging.basicConfig(
format='[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s',
level=logging.INFO,
handlers=[
logging.FileHandler(args.log_file, mode='w'),
logging.StreamHandler()
])
기본 로깅 설정을 구성하는데 basicConfig함수가 사용되는데, 이 함수는 로깅의 출력 형식, 대상, 레벨 등의 기본 설정을 간단하게 지정할 수 있도록 한다.
로깅 레벨은 총 5가지로 구성되는데, DEBUG, INFO, WARNING, ERROR, CRITICAL 등이 있다.
logging에서 handler는 위와 같은 로그들을 어디에 저장하거나 출력할때 사용하는데,
logging.FileHandler()의 경우 특정 파일에 로그를 저장시킬 수 있으며, logging.StreamHandler()의 경우 Command, Terminal 창에 로그를 띄울 수 있다.
optimizer = torch.optim.Adam([{
'params': pfld_backbone.parameters()
},{
'parmas': auxiliary.parameters()
}],
lr=args.base_lr,
weight_decay=args.weight_decay
)
optimizer를 사용할때, weight_decay를 사용하는 것을 볼 수 있는데, weight_decay는 무슨 동작을 하는 것일까?
Overfitting 문제를 해결하기 위해 사용한다. 우리가 모델을 학습한다고 가정해보자.
딥러닝의 대부분은 linear보단 nonlinear 형태로 학습을 진행한다. 대부분 두개 이상의 class를 갖기 때문이다. linear 형태를 한번 보자.
linear에서도 가중치가 큰 데이터가 있다면 직선이 y축 기준으로 살짝 올라가거나 내려가는 편향적인 모습을 보여주는데, nonlinear 데이터로 본다면 좀더 쉽게 와닿을 수 있다.

weigth_decay를 적용하지 않은 nonlinear 모델은 weight가 큰 데이터 쪽으로 크게 움직이는데, weight_decay를 적용한 nonlinear 모델은 weight가 큰 데이터 쪽으로 움직이지 않아 편향에서 멀어지는 방향으로 많은 도움이 된다. 이로써 overfitting을 방지할 수 있다는 것이고 local minimum에 빠질 가능성도 현저하게 줄어들게 된다.
Scheduler 사용 방법에도 여러가지가 있다.
특정 epoch에 다다르면 learning_rate를 감소시키는 방법이 있고 loss가 더이상 줄어들지 않으면 learning_rate를 더 줄이는 방법이 있다.
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', patience=args.lr_patience, verbose=True
)
mode를 min으로 설정할 경우 loss가 감소하지 않을 때 lr을 감소시키고 patience로 일정 epoch까지 기다렸다가 더이상 loss가 변하지 않으면 그때 lr을 감소시키기 만드는 방법도 있다.
if args.resume:
checkpoint = torch.load(args.resume)
auxiliary.load_state_dict(checkpoint["auxiliarynet"])
pfld_backbone.load_state_dict(checkpoint["pfld_backbone"])
args.start_epoch = checkpoint["epoch"]
resume sector는 추가적으로 더 학습을 시키고 싶을 경우 해당 모델을 불러와서 추가적인 학습을 행하기 위한 용도이다.
wandb.init(
project="pfld_pr",
name="pfld_face_landamrk1",
config={
"learning_rate": args.lr,
"epochs": args.end_epoch,
"batch_size": args.batch_size
}
)
for epoch in range(args.start_epoch, args.end_epoch + 1):
weighted_train_loss, train_loss = train(dataloader, pfld_backbone,
auxiliary, criterion,
optimizer, epoch)
filename = os.path.join(str(args.snapshot),
"checkpoint_epoch_" + str(epoch) + '.pth.tar')
save_checkpoint(
{
'epoch': epoch,
'pfld_backbone': pfld_backbone.state_dict(),
'auxiliarynet': auxiliary.state_dict()
}, filename)
val_loss = validate(wlfw_val_dataloader, pfld_backbone, auxiliary,
criterion)
scheduler.step(val_loss)
wandb.log({
'epoch': epoch,
'weighted_train_loss': weighted_train_loss,
'train_loss': train_loss,
'val_loss': val_loss
})
전 회사에서 tensorboard를 사용했었는데, W&B 사용으로 많은 점이 달라졌다. 우선 보기 편하다는 장점이 있고 깔끔하다. 충분히 tensorboard만으로 생산성, 실험 관리 등이 가능하지만 UI 적으로 사용자들에게 거부감없이 스무스하다(?)는 장점이 있는 것 같다.
전체 코드
import os
import cv2
import argparse
import logging
import torch
import wandb
import numpy as np
from torch.utils.data import DataLoader
from model import PFLDBackbone, AuxiliaryBlock
from dataset import WLFWDataset
from loss import PFLDloss
from utils import AverageMeter
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
def print_args(args):
for arg in vars(args):
s = arg + ': ' + str(getattr(args, arg))
logging.info(s)
def save_checkpoint(state, filename='checkpoint.pth.tar'):
torch.save(state, filename)
logging.info('Save checkpoint to {0:}'.format(filename))
def train(train_loader, pfld_backbone, auxiliary, criterion, optimizer):
losses = AverageMeter()
weighted_loss, loss = None, None
for img, landmark_gt, attribute_gt, euler_angle_gt in train_loader:
img = img.to(DEVICE)
landmark_gt = landmark_gt.to(DEVICE)
attribute_gt = attribute_gt.to(DEVICE)
euler_angle_gt = euler_angle_gt.to(DEVICE)
pfld_backbone = pfld_backbone.to(DEVICE)
auxiliary = auxiliary.to(DEVICE)
feature, landmarks = pfld_backbone(img)
euler_angle = auxiliary(feature)
weighted_loss, loss = criterion(attribute_gt, landmark_gt,
euler_angle_gt, euler_angle, landmarks, args.train_batchsize)
optimizer.zero_grad()
weighted_loss.backward()
optimizer.step()
losses.update(loss.item())
return weighted_loss, loss
def validate(valid_dataloader, pfld_backbone, auxiliary, criterion):
pfld_backbone.eval()
auxiliary.eval()
losses = []
with torch.no_grad():
for img, landmark_gt, attribute_gt, angle_gt in valid_dataloader:
img = img.to(DEVICE)
landmark_gt = landmark_gt.to(DEVICE)
attribute_gt = attribute_gt.to(DEVICE)
angle_gt = angle_gt.to(DEVICE)
pfld_backbone = pfld_backbone.to(DEVICE)
_, landmark = pfld_backbone(img)
loss = torch.mean(torch.sum((landmark_gt - landmark)**2, axis=1))
losses.append(loss.cpu().numpy())
print("====> Evalate:")
print('Eval set: Average loss: {:.4f}'.format(np.mean(losses)))
return np.mean(losses)
def main(args):
logging.basicConfig(
format=
'[%(asctime)s] [p%(process)s] [%(pathname)s:%(lineno)d] [%(levelname)s] %(message)s',
level=logging.INFO,
handlers=[
logging.FileHandler(args.log_file, mode='w'),
logging.StreamHandler()
])
print_args(args)
pfld_backbone = PFLDBackbone().to(DEVICE)
auxiliary = AuxiliaryBlock().to(DEVICE)
criterion = PFLDloss()
optimizer = torch.optim.Adam([{
'params': pfld_backbone.parameters()
},{
'parmas': auxiliary.parameters()
}],
lr=args.base_lr,
weight_decay=args.weight_decay
)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', patience=args.lr_patience, verbose=True
)
if args.resume:
checkpoint = torch.load(args.resume)
auxiliary.load_state_dict(checkpoint["auxiliarynet"])
pfld_backbone.load_state_dict(checkpoint["pfld_backbone"])
args.start_epoch = checkpoint["epoch"]
transform = transform.Compose([transform.ToTensor()])
wlfwdataset = WLFWDataset(args.dataroot, transform)
dataloader = DataLoader(wlfwdataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num.workers)
wlfw_val_dataset = WLFWDataset(args.val_dataroot,transform)
wlfw_val_dataloader = DataLoader(wlfw_val_dataset,
batch_size=args.val_batchsize,
shuffle=False,
num_workers=args.workers)
_
wandb.init(
project="pfld_pr",
name="pfld_face_landamrk1",
config={
"learning_rate": args.lr,
"epochs": args.end_epoch,
"batch_size": args.batch_size
}
)
for epoch in range(args.start_epoch, args.end_epoch + 1):
weighted_train_loss, train_loss = train(dataloader, pfld_backbone,
auxiliary, criterion,
optimizer, epoch)
filename = os.path.join(str(args.snapshot),
"checkpoint_epoch_" + str(epoch) + '.pth.tar')
save_checkpoint(
{
'epoch': epoch,
'pfld_backbone': pfld_backbone.state_dict(),
'auxiliarynet': auxiliary.state_dict()
}, filename)
val_loss = validate(wlfw_val_dataloader, pfld_backbone, auxiliary,
criterion)
scheduler.step(val_loss)
wandb.log({
'epoch': epoch,
'weighted_train_loss': weighted_train_loss,
'train_loss': train_loss,
'val_loss': val_loss
})
def parse_args():
parser = argparse.ArgumentParser(description="PFLD")
parser.add_argument('-j', '--workers', default=0, type=int)
parser.add_argument('--devices_id', default='0', type=str) #TBD
# parser.add_argument('--test_initial', default='false', type=str2bool) #TBD
# training
## -- optimizer
parser.add_argument('--base_lr', default=0.0001, type=int)
parser.add_argument('--weight-decay', '--wd', default=1e-6, type=float)
# -- lr
parser.add_argument("--lr_patience", default=40, type=int)
# -- epoch
parser.add_argument('--start_epoch', default=1, type=int)
parser.add_argument('--end_epoch', default=500, type=int)
# -- snapshot、tensorboard log and checkpoint
parser.add_argument('--snapshot',
default='./checkpoint/snapshot/',
type=str,
metavar='PATH')
parser.add_argument('--log_file',
default="./checkpoint/train.logs",
type=str)
parser.add_argument('--tensorboard',
default="./checkpoint/tensorboard",
type=str)
parser.add_argument(
'--resume',
default='',
type=str,
metavar='PATH')
# --dataset
parser.add_argument('--dataroot',
default='./data/train_data/list.txt',
type=str,
metavar='PATH')
parser.add_argument('--val_dataroot',
default='./data/test_data/list.txt',
type=str,
metavar='PATH')
parser.add_argument('--train_batchsize', default=256, type=int)
parser.add_argument('--val_batchsize', default=256, type=int)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
main(args)