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)