def test_compare_crop_and_pad(img_dtype, px, percent, pad_mode, pad_cval, keep_size):
h, w, c = 100, 100, 3
mode_mapping = {
cv2.BORDER_CONSTANT: "constant",
cv2.BORDER_REPLICATE: "edge",
cv2.BORDER_REFLECT101: "reflect",
cv2.BORDER_WRAP: "wrap",
}
pad_mode_iaa = mode_mapping[pad_mode]
bbox_params = A.BboxParams(format="pascal_voc")
keypoint_params = A.KeypointParams(format="xy", remove_invisible=False)
keypoints = np.random.randint(0, min(h, w), [10, 2])
bboxes = []
for i in range(10):
x1, y1 = np.random.randint(0, min(h, w) - 2, 2)
x2 = np.random.randint(x1 + 1, w - 1)
y2 = np.random.randint(y1 + 1, h - 1)
bboxes.append([x1, y1, x2, y2, 0])
transform_albu = A.Compose(
[
A.CropAndPad(
px=px,
percent=percent,
pad_mode=pad_mode,
pad_cval=pad_cval,
keep_size=keep_size,
p=1,
interpolation=cv2.INTER_AREA
if (px is not None and px < 0) or (percent is not None and percent < 0)
else cv2.INTER_LINEAR,
)
],
bbox_params=bbox_params,
keypoint_params=keypoint_params,
)
transform_iaa = A.Compose(
[A.IAACropAndPad(px=px, percent=percent, pad_mode=pad_mode_iaa, pad_cval=pad_cval, keep_size=keep_size, p=1)],
bbox_params=bbox_params,
keypoint_params=keypoint_params,
)
if img_dtype == np.uint8:
img = np.random.randint(0, 256, (h, w, c), dtype=np.uint8)
else:
img = np.random.random((h, w, c)).astype(img_dtype)
res_albu = transform_albu(image=img, keypoints=keypoints, bboxes=bboxes)
res_iaa = transform_iaa(image=img, keypoints=keypoints, bboxes=bboxes)
for key, item in res_albu.items():
if key == "bboxes":
bboxes = np.array(res_iaa[key])
h = bboxes[:, 3] - bboxes[:, 1]
w = bboxes[:, 2] - bboxes[:, 0]
res_iaa[key] = bboxes[(h > 0) & (w > 0)]
assert np.allclose(item, res_iaa[key]), f"{key} are not equal"
def main():
# 주요 path 정의
data_path = './data'
train_dir = Path(data_path, 'images/train_imgs')
# config 파일을 가져옵니다.
args = parse_args()
update_config(cfg, args)
lr = cfg.TRAIN.LR
lamb = cfg.LAMB
test_option = eval(cfg.test_option)
input_w = cfg.MODEL.IMAGE_SIZE[1]
input_h = cfg.MODEL.IMAGE_SIZE[0]
# 랜덤 요소를 최대한 줄여줌
RANDOM_SEED = int(cfg.RANDOMSEED)
np.random.seed(RANDOM_SEED) # cpu vars
torch.manual_seed(RANDOM_SEED) # cpu vars
random.seed(RANDOM_SEED) # Python
os.environ['PYTHONHASHSEED'] = str(RANDOM_SEED) # Python hash buildin
torch.backends.cudnn.deterministic = True #needed
torch.backends.cudnn.benchmark = False
torch.cuda.manual_seed(RANDOM_SEED)
torch.cuda.manual_seed_all(RANDOM_SEED) # if use multi-GPU
# log 데이터와 최종 저장위치를 만듭니다.
logger, final_output_dir, tb_log_dir = create_logger(cfg, args.cfg, f'lr_{str(lr)}', 'train')
logger.info(pprint.pformat(args))
logger.info(cfg)
# cudnn related setting
cudnn.benchmark = cfg.CUDNN.BENCHMARK
# annotation 파일을 만듭니다.
if os.path.isfile(data_path+'/annotations/train_annotation.pkl') == False :
make_annotations(data_path)
# 쓰려는 모델을 불러옵니다.
model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
cfg, is_train=True
)
# model의 끝부분 수정 및 초기화 작업을 진행합니다.
model = initialize_model(model, cfg)
# model 파일과 train.py 파일을 copy합니다.
this_dir = os.path.dirname(__file__)
shutil.copy2(
os.path.join(this_dir, '../lib/models', cfg.MODEL.NAME + '.py'),
final_output_dir)
shutil.copy2(
os.path.join(this_dir, '../tools', 'train.py'),
final_output_dir)
writer_dict = {
'writer': SummaryWriter(log_dir=tb_log_dir),
'train_global_steps': 0,
'valid_global_steps': 0,
}
# model을 그래픽카드가 있을 경우 cuda device로 전환합니다.
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
# loss를 정의합니다.
criterion = nn.MSELoss().cuda()
# Data Augumentation을 정의합니다.
A_transforms = {
'val':
A.Compose([
A.Resize(input_h, input_w, always_apply=True),
A.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
ToTensorV2()
], bbox_params=A.BboxParams(format="coco", min_visibility=0.05, label_fields=['class_labels'])),
'test':
A.Compose([
A.Resize(input_h, input_w, always_apply=True),
A.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
ToTensorV2()
])
}
if input_h == input_w :
A_transforms['train'] = A.Compose([
A.Resize(input_h, input_w, always_apply=True),
A.OneOf([A.HorizontalFlip(p=1),
A.VerticalFlip(p=1),
A.Rotate(p=1),
A.RandomRotate90(p=1)
], p=0.5),
A.OneOf([A.MotionBlur(p=1),
A.GaussNoise(p=1),
A.ColorJitter(p=1)
], p=0.5),
A.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
ToTensorV2()
], bbox_params=A.BboxParams(format="coco", min_visibility=0.05, label_fields=['class_labels']))
else :
A_transforms['train'] = A.Compose([
A.Resize(input_h, input_w, always_apply=True),
A.OneOf([A.HorizontalFlip(p=1),
A.VerticalFlip(p=1),
A.Rotate(p=1),
], p=0.5),
A.OneOf([A.MotionBlur(p=1),
A.GaussNoise(p=1)
], p=0.5),
A.OneOf([A.CropAndPad(percent=0.1, p=1),
A.CropAndPad(percent=0.2, p=1),
A.CropAndPad(percent=0.3, p=1)
], p=0.5),
A.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
ToTensorV2()
], bbox_params=A.BboxParams(format="coco", min_visibility=0.05, label_fields=['class_labels']))
# parameter를 설정합니다.
batch_size = int(cfg.TRAIN.BATCH_SIZE_PER_GPU)
test_ratio = float(cfg.TEST_RATIO)
num_epochs = cfg.TRAIN.END_EPOCH
# earlystopping에 주는 숫자 변수입니다.
num_earlystop = num_epochs
# torch에서 사용할 dataset을 생성합니다.
imgs, bbox, class_labels = make_train_data(data_path)
since = time.time()
"""
# test_option : train, valid로 데이터를 나눌 때 test data를 고려할지 결정합니다.
* True일 경우 test file을 10% 뺍니다.
* False일 경우 test file 빼지 않습니다.
"""
if test_option == True :
X_train, X_test, y_train, y_test = train_test_split(imgs, bbox, test_size=0.1, random_state=RANDOM_SEED)
test_dataset = [X_test, y_test]
with open(final_output_dir+'/test_dataset.pkl', 'wb') as f:
pickle.dump(test_dataset, f)
X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=test_ratio, random_state=RANDOM_SEED)
test_data = Dataset(train_dir, X_test, y_test, data_transforms=A_transforms, class_labels=class_labels, phase='val')
test_loader = data_utils.DataLoader(test_data, batch_size=batch_size, shuffle=False)
else :
X_train, X_val, y_train, y_val = train_test_split(imgs, bbox, test_size=test_ratio, random_state=RANDOM_SEED)
train_data = Dataset(train_dir, X_train, y_train, data_transforms=A_transforms, class_labels=class_labels, phase='train')
val_data = Dataset(train_dir, X_val, y_val, data_transforms=A_transforms, class_labels=class_labels, phase='val')
train_loader = data_utils.DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = data_utils.DataLoader(val_data, batch_size=batch_size, shuffle=False)
# best loss를 판별하기 위한 변수 초기화
best_perf = 10000000000
test_loss = None
best_model = False
# optimizer 정의
optimizer = optim.Adam(
model.parameters(),
lr=lr
)
# 중간에 학습된 모델이 있다면 해당 epoch에서부터 진행할 수 있도록 만듭니다.
begin_epoch = cfg.TRAIN.BEGIN_EPOCH
checkpoint_file = os.path.join(
final_output_dir, 'checkpoint.pth'
)
if cfg.AUTO_RESUME and os.path.exists(checkpoint_file):
logger.info("=> loading checkpoint '{}'".format(checkpoint_file))
checkpoint = torch.load(checkpoint_file)
begin_epoch = checkpoint['epoch']
best_perf = checkpoint['perf']
num_epochs = checkpoint['epoch']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
logger.info("=> loaded checkpoint '{}' (epoch {})".format(
checkpoint_file, checkpoint['epoch']))
# lr_scheduler 정의
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, cfg.TRAIN.LR_STEP, cfg.TRAIN.LR_FACTOR,
last_epoch=-1
)
# early stopping하는데 사용하는 count 변수
count = 0
val_losses = []
train_losses = []
# 학습 시작
for epoch in range(begin_epoch, num_epochs):
epoch_since = time.time()
lr_scheduler.step()
# train for one epoch
train_loss = train(cfg, device, train_loader, model, criterion, optimizer, epoch,
final_output_dir, tb_log_dir, writer_dict, lamb=lamb)
# evaluate on validation set
perf_indicator = validate(
cfg, device, val_loader, val_data, model, criterion,
final_output_dir, tb_log_dir, writer_dict, lamb=lamb
)
# 해당 epoch이 best_model인지 판별합니다. valid 값을 기준으로 결정됩니다.
if perf_indicator <= best_perf:
best_perf = perf_indicator
best_model = True
count = 0
else:
best_model = False
count +=1
logger.info('=> saving checkpoint to {}'.format(final_output_dir))
save_checkpoint({
'epoch': epoch + 1,
'model': cfg.MODEL.NAME,
'state_dict': model.state_dict(),
'best_state_dict': model.state_dict(),
'perf': perf_indicator,
'optimizer': optimizer.state_dict(),
}, best_model, final_output_dir)
# loss를 저장합니다.
val_losses.append(perf_indicator)
train_losses.append(train_loss)
if count == num_earlystop :
break
epoch_time_elapsed = time.time() - epoch_since
print(f'epoch : {epoch}' \
f' train loss : {round(train_loss,3)}' \
f' valid loss : {round(perf_indicator,3)}' \
f' Elapsed time: {int(epoch_time_elapsed // 60)}m {int(epoch_time_elapsed % 60)}s')
# log 파일 등을 저장합니다.
final_model_state_file = os.path.join(
final_output_dir, 'final_state.pth'
)
logger.info('=> saving final model state to {}'.format(
final_model_state_file)
)
torch.save(model.state_dict(), final_model_state_file)
writer_dict['writer'].close()
time_elapsed = time.time() - since
print('Training and Validation complete in {:.0f}m {:.0f}s'.format(time_elapsed // 60, time_elapsed % 60))
print('Best validation loss: {:4f}\n'.format(best_perf))
# test_option이 True일 경우, 떼어난 10% 데이터에 대해 만들어진 모델로 eval을 진행합니다.
if test_option == True :
# test data
model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
cfg, is_train=True)
model = initialize_model(model, cfg)
parameters = f'{final_output_dir}/model_best.pth'
model = model.to(device)
model.load_state_dict(torch.load(parameters))
test_loss = validate(
cfg, device, test_loader, test_data, model, criterion,
final_output_dir, tb_log_dir, writer_dict, lamb=lamb
)
print(f'test loss : {test_loss}')
# loss 결과를 pickle 파일로 따로 저장합니다.
result_dict = {}
result_dict['val_loss'] = val_losses
result_dict['train_loss'] = train_losses
result_dict['best_loss'] = best_perf
result_dict['test_loss'] = test_loss
result_dict['lr'] = lr
with open(final_output_dir+'/result.pkl', 'wb') as f:
pickle.dump(result_dict, f)