def main(): """Parameters initialization and starting attention model training """ # read command line arguments args = get_parser().parse_args() # set random seed seed_everything(args.seed) # paths to dataset train_path = osp.join(args.dataset_path, 'train') test_path = osp.join(args.dataset_path, 'test') # declare Unet model with two ouput classes (occluders and their shadows) model = smp.Unet(encoder_name=args.encoder, classes=2, activation='sigmoid',) # replace the first convolutional layer in model: 4 channels tensor as model input model.encoder.conv1 = nn.Conv2d(4, 64, kernel_size=(7, 7), stride=(2, 2), \ padding=(3, 3), bias=False) # declare datasets train_dataset = ARDataset(train_path, augmentation=get_training_augmentation(args.img_size), \ preprocessing=get_preprocessing(),) valid_dataset = ARDataset(test_path, augmentation=get_validation_augmentation(args.img_size), \ preprocessing=get_preprocessing(),) # declare loaders train_loader = DataLoader(train_dataset, batch_size=args.batch_size, \ shuffle=True, num_workers=args.num_workers) valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, \ shuffle=False, num_workers=args.num_workers) # declare loss function, optimizer and metric loss = smp.utils.losses.DiceLoss() metric = smp.utils.metrics.IoU(threshold=args.iou_th) optimizer = torch.optim.Adam([dict(params=model.parameters(), lr=args.lr),]) # tensorboard writer = SummaryWriter() # device device = torch.device(args.device if torch.cuda.is_available() else 'cpu') # start training train( writer=writer, n_epoch=args.n_epoch, train_loader=train_loader, valid_loader=valid_loader, model_path=args.model_path, model=model, loss=loss, metric=metric, optimizer=optimizer, device=device )
def get_train_val_loaders(encoder_type, batch_size=16, pseudo_label=False): if encoder_type.startswith('myunet'): encoder_type = 'resnet50' preprocessing_fn = smp.encoders.get_preprocessing_fn( encoder_type, 'imagenet') train, train_ids, valid_ids = prepare_df() train['pseudo'] = 0 pseudo_imgs = set() if pseudo_label: train_pseudo = prepare_df( train_file=f'{settings.DATA_DIR}/sub_blend_1111_1.csv', pseudo_label=True) train_pseudo['pseudo'] = 1 pseudo_ids = train_pseudo.im_id.unique().tolist() print(pseudo_ids[:10]) pseudo_imgs = set(pseudo_ids) train_ids.extend(pseudo_ids) train = pd.concat([train, train_pseudo]) print(train.head()) print(train_pseudo.head()) print(train.shape) print(len(train_ids)) num_workers = 24 train_dataset = CloudDataset( df=train, datatype='train', img_ids=train_ids, transforms=get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), pseudo_imgs=pseudo_imgs) valid_dataset = CloudDataset( df=train, datatype='valid', img_ids=valid_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers) valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers) train_loader.num = len(train_ids) valid_loader.num = len(valid_ids) loaders = {"train": train_loader, "valid": valid_loader} return loaders
def get_test_loader(encoder_type, batch_size=16): if encoder_type.startswith('myunet'): encoder_type = 'resnet50' preprocessing_fn = smp.encoders.get_preprocessing_fn( encoder_type, 'imagenet') sub = pd.read_csv(os.path.join(settings.DATA_DIR, 'sample_submission.csv')) sub['label'] = sub['Image_Label'].apply(lambda x: x.split('_')[1]) sub['im_id'] = sub['Image_Label'].apply(lambda x: x.split('_')[0]) test_ids = sub['Image_Label'].apply( lambda x: x.split('_')[0]).drop_duplicates().values test_dataset = CloudDataset( df=sub, datatype='test', img_ids=test_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=24) return test_loader
def Predict(self, img_path, vis=True): ''' User function: Run inference on image and visualize it. Output mask saved as output_mask.npy Args: img_path (str): Relative path to the image file vis (bool): If True, predicted mask is displayed. Returns: list: List of bounding box locations of predicted objects along with classes. ''' dirPath = "tmp_test" if (os.path.isdir(dirPath)): shutil.rmtree(dirPath) os.mkdir(dirPath) os.mkdir(dirPath + "/img_dir") os.mkdir(dirPath + "/gt_dir") os.system("cp " + img_path + " " + dirPath + "/img_dir") os.system("cp " + img_path + " " + dirPath + "/gt_dir") x_test_dir = dirPath + "/img_dir" y_test_dir = dirPath + "/gt_dir" if (self.system_dict["params"]["image_shape"][0] % 32 != 0): self.system_dict["params"]["image_shape"][0] += ( 32 - self.system_dict["params"]["image_shape"][0] % 32) if (self.system_dict["params"]["image_shape"][1] % 32 != 0): self.system_dict["params"]["image_shape"][1] += ( 32 - self.system_dict["params"]["image_shape"][1] % 32) preprocess_input = sm.get_preprocessing( self.system_dict["params"]["backbone"]) test_dataset = Dataset( x_test_dir, y_test_dir, self.system_dict["params"]["classes_dict"], classes_to_train=self.system_dict["params"]["classes_to_train"], augmentation=get_validation_augmentation( self.system_dict["params"]["image_shape"][0], self.system_dict["params"]["image_shape"][1]), preprocessing=get_preprocessing(preprocess_input), ) test_dataloader = Dataloder(test_dataset, batch_size=1, shuffle=False) image, gt_mask = test_dataset[0] image = np.expand_dims(image, axis=0) pr_mask = self.system_dict["local"]["model"].predict(image).round() np.save("output_mask.npy", pr_mask) if (vis): visualize( image=denormalize(image.squeeze()), pr_mask=pr_mask[..., 0].squeeze(), )
def main(args): """ Main code for creating the segmentation-only submission file. All masks are converted to either "" or RLEs Args: args (instance of argparse.ArgumentParser): arguments must be compiled with parse_args Returns: None """ torch.cuda.empty_cache() gc.collect() attention_type = None if args.attention_type == "None" else args.attention_type model = smp.Unet(encoder_name=args.encoder, encoder_weights=None, classes=4, activation=None, attention_type=attention_type) # setting up the test I/O preprocessing_fn = smp.encoders.get_preprocessing_fn( args.encoder, "imagenet") # setting up the train/val split with filenames train, sub, _ = setup_train_and_sub_df(args.dset_path) test_ids = sub["ImageId_ClassId"].apply( lambda x: x.split("_")[0]).drop_duplicates().values # datasets/data loaders test_dataset = SteelDataset( args.dset_path, df=sub, datatype="test", im_ids=test_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) loaders = {"test": test_loader} # loading the pickled class_params if they exist class_params_path = os.path.join(args.dset_path, "class_params.pickle") if os.path.exists(class_params_path): print(f"Loading {class_params_path}...") # Load data (deserialize) with open(class_params_path, "rb") as handle: class_params = pickle.load(handle) else: class_params = "default" create_submission(args.checkpoint_path, model=model, loaders=loaders, sub=sub, class_params=class_params)
def test_arshadowgan(): """Test ARShadowGAN-like architecture """ # parse command line arguments args = get_parser().parse_args() # create folders if not osp.exists(args.result_path): os.makedirs(args.result_path) # dataset and dataloader declaration dataset = ARDataset(args.dataset_path, augmentation=get_validation_augmentation( args.img_size), preprocessing=get_preprocessing(), is_train=False) dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) # define device device = torch.device(args.device if torch.cuda.is_available() else 'cpu') # define model model = ARShadowGAN(encoder_att=args.attention_encoder, encoder_SG=args.SG_encoder, model_path_attention=args.path_att, model_path_SG=args.path_SG, device=device) # model in eval mode now model.eval() # inference counter = 0 for i, data in enumerate(tqdm(dataloader)): tensor_att = torch.cat( (data[0][:, :3], torch.unsqueeze(data[1][:, -1], axis=1)), axis=1).to(device) tensor_SG = torch.cat( (data[2][:, :3], torch.unsqueeze(data[3][:, -1], axis=1)), axis=1).to(device) with torch.no_grad(): result, output_mask1 = model(tensor_att, tensor_SG) for j in range(args.batch_size): counter += 1 output_image = np.uint8( 127.5 * (result.cpu().numpy()[j].transpose(1, 2, 0) + 1.0)) output_image = cv2.cvtColor(output_image, cv2.COLOR_BGR2RGB) cv2.imwrite(osp.join(args.result_path, str(counter) + '.png'), output_image)
def main(args): """ Main code for creating the classification submission file. No masks predictions will be blank, and predictions for masks will be "1". Args: args (instance of argparse.ArgumentParser): arguments must be compiled with parse_args Returns: None """ torch.cuda.empty_cache() gc.collect() model = ResNet34(pre=None, num_classes=4, use_simple_head=True) # setting up the test I/O preprocessing_fn = smp.encoders.get_preprocessing_fn("resnet34", "imagenet") train, sub, _ = setup_train_and_sub_df(args.dset_path) test_ids = sub["ImageId_ClassId"].apply(lambda x: x.split("_")[0]).drop_duplicates().values # datasets/data loaders test_dataset = ClassificationSteelDataset( args.dset_path, df=sub, datatype="test", im_ids=test_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn) ) test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) loaders = {"test": test_loader} # loading the pickled class_params if they exist class_params_path = os.path.join(args.dset_path, "class_params_classification.pickle") if os.path.exists(class_params_path): print(f"Loading {class_params_path}...") # Load data (deserialize) with open(class_params_path, "rb") as handle: class_params = pickle.load(handle) else: class_params = "default" create_submission(args.checkpoint_path, model=model, loaders=loaders, sub=sub, class_params=class_params)
def get_train_val_loaders(encoder_type, batch_size=16, ifold=0): if encoder_type.startswith('myunet'): encoder_type = 'resnet50' preprocessing_fn = smp.encoders.get_preprocessing_fn(encoder_type, 'imagenet') train, train_ids, valid_ids = prepare_df(ifold=ifold) print('val:', valid_ids[:10]) num_workers = 24 train_dataset = CloudDataset(df=train, datatype='train', img_ids=train_ids, transforms = get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) valid_dataset = CloudDataset(df=train, datatype='valid', img_ids=valid_ids, transforms = get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn)) train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers) valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers) train_loader.num = len(train_ids) valid_loader.num = len(valid_ids) loaders = { "train": train_loader, "valid": valid_loader } return loaders
def main(args): """ Main code for training for training a U-Net with some user-defined encoder. Args: args (instance of argparse.ArgumentParser): arguments must be compiled with parse_args Returns: None """ # setting up the train/val split with filenames train, sub, id_mask_count = setup_train_and_sub_df(args.dset_path) # setting up the train/val split with filenames seed_everything(args.split_seed) train_ids, valid_ids = train_test_split(id_mask_count["im_id"].values, random_state=args.split_seed, stratify=id_mask_count["count"], test_size=args.test_size) # setting up model (U-Net with ImageNet Encoders) ENCODER_WEIGHTS = "imagenet" DEVICE = "cuda" attention_type = None if args.attention_type == "None" else args.attention_type model = smp.Unet(encoder_name=args.encoder, encoder_weights=ENCODER_WEIGHTS, classes=4, activation=None, attention_type=attention_type) preprocessing_fn = smp.encoders.get_preprocessing_fn( args.encoder, ENCODER_WEIGHTS) # Setting up the I/O train_dataset = SteelDataset( args.dset_path, df=train, datatype="train", im_ids=train_ids, transforms=get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), use_resized_dataset=args.use_resized_dataset) valid_dataset = SteelDataset( args.dset_path, df=train, datatype="valid", im_ids=valid_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), use_resized_dataset=args.use_resized_dataset) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) loaders = {"train": train_loader, "valid": valid_loader} # everything is saved here (i.e. weights + stats) logdir = "./logs/segmentation" # model, criterion, optimizer optimizer = torch.optim.Adam([ { "params": model.decoder.parameters(), "lr": args.encoder_lr }, { "params": model.encoder.parameters(), "lr": args.decoder_lr }, ]) scheduler = ReduceLROnPlateau(optimizer, factor=0.15, patience=2) criterion = smp.utils.losses.BCEDiceLoss(eps=1.) runner = SupervisedRunner() callbacks_list = [ DiceCallback(), EarlyStoppingCallback(patience=5, min_delta=0.001), ] if args.checkpoint_path != "None": # hacky way to say no checkpoint callback but eh what the heck ckpoint_p = Path(args.checkpoint_path) fname = ckpoint_p.name resume_dir = str(ckpoint_p.parents[0] ) # everything in the path besides the base file name print( f"Loading {fname} from {resume_dir}. Checkpoints will also be saved in {resume_dir}." ) callbacks_list = callbacks_list + [ CheckpointCallback(resume=fname, resume_dir=resume_dir), ] runner.train(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, loaders=loaders, callbacks=callbacks_list, logdir=logdir, num_epochs=args.num_epochs, verbose=True)
def main(args): """ Main code for training a classification model. Args: args (instance of argparse.ArgumentParser): arguments must be compiled with parse_args Returns: None """ # Reading the in the .csvs train = pd.read_csv(os.path.join(args.dset_path, "train.csv")) sub = pd.read_csv(os.path.join(args.dset_path, "sample_submission.csv")) # setting up the train/val split with filenames train, sub, id_mask_count = setup_train_and_sub_df(args.dset_path) # setting up the train/val split with filenames seed_everything(args.split_seed) train_ids, valid_ids = train_test_split(id_mask_count["im_id"].values, random_state=args.split_seed, stratify=id_mask_count["count"], test_size=args.test_size) # setting up the classification model ENCODER_WEIGHTS = "imagenet" DEVICE = "cuda" model = ResNet34(pre=ENCODER_WEIGHTS, num_classes=4, use_simple_head=True) preprocessing_fn = smp.encoders.get_preprocessing_fn( "resnet34", ENCODER_WEIGHTS) # Setting up the I/O train_dataset = ClassificationSteelDataset( args.dset_path, df=train, datatype="train", im_ids=train_ids, transforms=get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), ) valid_dataset = ClassificationSteelDataset( args.dset_path, df=train, datatype="valid", im_ids=valid_ids, transforms=get_validation_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), ) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) loaders = {"train": train_loader, "valid": valid_loader} # everything is saved here (i.e. weights + stats) logdir = "./logs/segmentation" # model, criterion, optimizer optimizer = torch.optim.Adam(model.parameters(), lr=3e-4) scheduler = ReduceLROnPlateau(optimizer, factor=0.15, patience=2) criterion = smp.utils.losses.BCEDiceLoss(eps=1.) runner = SupervisedRunner() runner.train(model=model, criterion=criterion, optimizer=optimizer, scheduler=scheduler, loaders=loaders, callbacks=[ DiceCallback(), EarlyStoppingCallback(patience=5, min_delta=0.001) ], logdir=logdir, num_epochs=args.num_epochs, verbose=True) utils.plot_metrics( logdir=logdir, # specify which metrics we want to plot metrics=["loss", "dice", "lr", "_base/lr"])
reset_index().rename(columns={'index': 'img_id', 'Image_Label': 'count'}) train_ids, valid_ids = train_test_split(id_mask_count['img_id'].values, random_state=42, stratify=id_mask_count['count'], test_size=0.1) test_ids = sub['Image_Label'].apply(lambda x: x.split('_')[0]).drop_duplicates().values print("creating preprocessing module...") ENCODER = 'resnet50' ENCODER_WEIGHTS = 'imagenet' preprocessing_fn = smp.encoders.get_preprocessing_fn(ENCODER, ENCODER_WEIGHTS) print("creating data loader...") num_workers = 4 bs = 16 train_dataset = CloudDataset(path = path, df=train, datatype='train', img_ids=train_ids, transforms = utils.get_training_augmentation(), preprocessing = utils.get_preprocessing(preprocessing_fn)) valid_dataset = CloudDataset(path = path, df=train, datatype='valid', img_ids=valid_ids, transforms = utils.get_validation_augmentation(), preprocessing = utils.get_preprocessing(preprocessing_fn)) train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=num_workers) valid_loader = DataLoader(valid_dataset, batch_size=bs, shuffle=False, num_workers=num_workers) loaders = { "train": train_loader, "valid": valid_loader } print("setting for training...") ACTIVATION = None model = smp.DeepLabV3Plus( encoder_name=ENCODER,
def main(config): # load best saved checkpoint best_model = torch.load(config.TRAINED_MODEL) # send to cuda and change mode `eval` best_model = best_model.to(config.DEVICE) best_model.eval() # create Dataset and DataLoader preprocessing_fn = normalization # set augmentation test_augmentation = get_augmentation_wrapper(config.TEST_AUGORATION) # predict each videos for test_dir in config.TEST_DIRS: # set dataset test_dataset = VideoPredictDataset( os.path.join(config.DATA_DIR, test_dir), augmentation=test_augmentation(height=config.HEIGHT, width=config.WIDTH), preprocessing=get_preprocessing(preprocessing_fn), skip_frame=config.SKIP_FRAME, ) video_name = os.path.basename(test_dataset.video_path) print('test on', video_name) print( f'number of frame: {len(test_dataset)} / {test_dataset.num_frames}' ) # predict per frame results = [] for frame_id, image in tqdm(test_dataset): # cast to tensor x_tensor = torch.from_numpy(image[np.newaxis]).to( config.DEVICE) # (1,C,H,W) # inference with torch.no_grad(): prediction = best_model.forward(x_tensor) # (1,C) # cast to ndarray prediction = prediction.cpu().numpy().round() # cast from onehot to class_id predicted_label = np.argmax(prediction, axis=1).squeeze() # (1,) # if apply `skip_frame`, padding previous label for i in range(config.SKIP_FRAME): # if last frame of video is over, stop padding if frame_id + 1 + i > test_dataset.num_frames: break results.append([frame_id + 1 + i, predicted_label]) # 答えのframe_idが1から始まるのでずらす # save prediction as csv save_dir = os.path.join(config.RESULT_DIR, test_dir) os.makedirs(save_dir, exist_ok=True) with open(os.path.join(save_dir, video_name.replace('.mp4', '.csv')), 'w') as fw: writer = csv.writer(fw) writer.writerow(['Frame', 'Steps']) # writer header writer.writerows(results)
train_csv = os.path.join(DATA_PATH,'train.csv') kfold_path = 'kfold.pkl' data_path = '/home/jianglb/pythonproject/cloud_segment/data/train_{}_{}'.format(*RESIZE) test_data = os.path.join(DATA_PATH,'test_{}_{}'.format(*RESIZE)) encoder_weights = 'instagram' tg_masks = [] tg_classes = [] tg_names = [] with torch.no_grad(): for fold in range(K): # validate preprocessing_fn = smp.encoders.get_preprocessing_fn('resnext101_32x8d', encoder_weights) validate_dataset = CloudTrainDataset2(train_csv, data_path, kfold_path, fold, phase='validate', transform_type=0, preprocessing=utils.get_preprocessing(preprocessing_fn),) validate_dataloader = DataLoader(validate_dataset,batch_size=BATCH_SIZE,num_workers=4,shuffle=False) for i, (names, imgs, masks, classes) in enumerate(tqdm(validate_dataloader)): tg_masks.append(np.around(masks.numpy() * 100).astype(np.uint8)) tg_classes.append(classes.numpy()) tg_names.extend(names) # tg_masks = np.concatenate(tg_masks,axis=0).astype(np.uint8) # tg_classes = np.concatenate(tg_classes,axis=0) # print(cloud_mask.max(),cloud_mask.min()) del validate_dataset
def main(): """Parameters initialization and starting SG model training """ # read command line arguments args = get_parser().parse_args() # set random seed seed_everything(args.seed) # paths to dataset train_path = osp.join(args.dataset_path, 'train') test_path = osp.join(args.dataset_path, 'test') # declare generator and discriminator models generator = Generator_with_Refin(args.encoder) discriminator = Discriminator(input_shape=(3,args.img_size,args.img_size)) # load weights if args.gen_weights != '': generator.load_state_dict(torch.load(args.gen_weights)) print('Generator weights loaded!') if args.discr_weights != '': discriminator.load_state_dict(torch.load(args.discr_weights)) print('Discriminator weights loaded!') # declare datasets train_dataset = ARDataset(train_path, augmentation=get_training_augmentation(args.img_size), augmentation_images=get_image_augmentation(), preprocessing=get_preprocessing(),) valid_dataset = ARDataset(test_path, augmentation=get_validation_augmentation(args.img_size), preprocessing=get_preprocessing(),) # declare loaders train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) # declare loss functions, optimizers and scheduler l2loss = nn.MSELoss() perloss = ContentLoss(feature_extractor="vgg16", layers=("relu3_3", )) GANloss = nn.MSELoss() optimizer_G = torch.optim.Adam([dict(params=generator.parameters(), lr=args.lr_G),]) optimizer_D = torch.optim.Adam([dict(params=discriminator.parameters(), lr=args.lr_D),]) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer_G, mode='min', factor=0.9, patience=args.patience) # device device = torch.device(args.device if torch.cuda.is_available() else 'cpu') # tensorboard writer = SummaryWriter() # start training train( generator=generator, discriminator=discriminator, device=device, n_epoch=args.n_epoch, optimizer_G=optimizer_G, optimizer_D=optimizer_D, train_loader=train_loader, valid_loader=valid_loader, scheduler=scheduler, losses=[l2loss, perloss, GANloss], models_paths=[args.Gmodel_path, args.Dmodel_path], bettas=[args.betta1, args.betta2, args.betta3], writer=writer, )
def main(config): ### model settings ### # create segmentation model with pretrained encoder model = getattr(models, config.MODEL)(num_classes=config.N_CLASSES, # pretrained=False, ) if config.MULTI: model = torch.nn.DataParallel(model) print('training with multi GPU') # create loss function loss = getattr(losses, config.LOSS)(**config.loss_params) # set metrics(except `Idle` class) metrics = [ utils.metrics.Accuracy(ignore_channels=[0]), utils.metrics.Precision(ignore_channels=[0]), utils.metrics.Recall(ignore_channels=[0]), utils.metrics.Fscore(ignore_channels=[0]), ] # set optimizer optimizer = getattr(torch.optim, config.OPTIMIZER)( params=model.parameters(), lr=config.LR, **config.optim_params, ) # learning rate scheduler scheduler = getattr(torch.optim.lr_scheduler, config.SCHEDULER)(optimizer, **config.scheduler_params) ### data provider settings ### # set dataset path x_train_dirs = [ os.path.join(config.DATA_DIR, train_dir) for train_dir in config.TRAINING_DIRS ] x_valid_dirs = [ os.path.join(config.DATA_DIR, train_dir) for train_dir in config.VALIDATION_DIRS ] # set augmentation training_augmentation = get_augmentation_wrapper( config.TRAINING_AUGORATION) validation_augmentation = get_augmentation_wrapper( config.VALIDATION_AUGORATION) # create Dataset and DataLoader preprocessing_fn = normalization train_dataset = ClassificationVideoDataset( x_train_dirs, augmentation=training_augmentation(height=config.HEIGHT, width=config.WIDTH), preprocessing=get_preprocessing(preprocessing_fn), classes=config.CLASSES, skip_frame=config.SKIP_FRAME, ) config.NUM_TRAINING = len(train_dataset) print('number of train data:', len(train_dataset)) valid_dataset = ClassificationVideoDataset( x_valid_dirs, augmentation=validation_augmentation(height=config.HEIGHT, width=config.WIDTH), preprocessing=get_preprocessing(preprocessing_fn), classes=config.CLASSES, skip_frame=config.SKIP_FRAME, ) config.NUM_VALIDATION = len(valid_dataset) print('number of validation data:', len(valid_dataset)) # overwrite settings yaml for update number of data for training and validation with open(os.path.join(config.RESULT_DIR, 'parameters.yml'), 'w') as fw: fw.write(yaml.dump(asdict(config))) train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE, shuffle=False, num_workers=12, drop_last=True) valid_loader = DataLoader( valid_dataset, batch_size=1, shuffle=False, num_workers=4, ) ### Trainer ### # create epoch runners train_epoch = utils.train.TrainEpoch( model, loss=loss, metrics=metrics, optimizer=optimizer, device=config.DEVICE, verbose=True, ) valid_epoch = utils.train.ValidEpoch( model, loss=loss, metrics=metrics, device=config.DEVICE, verbose=True, ) # train model loggers = { 'epoch': [], 'loss': [], 'val_loss': [], 'accuracy': [], 'val_accuracy': [], 'precision': [], 'val_precision': [], 'recall': [], 'val_recall': [], 'fscore': [], 'val_fscore': [], } # for callbacks max_score = 0 for i in range(0, config.EPOCHS): print('\nEpoch: {}'.format(i + 1)) train_logs = train_epoch.run(train_loader) valid_logs = valid_epoch.run(valid_loader) ## logger ## # TODO: 冗長すぎる loggers['epoch'].append(i + 1) loggers['loss'].append(train_logs['loss']) loggers['val_loss'].append(valid_logs['loss']) loggers['accuracy'].append(train_logs['accuracy']) loggers['val_accuracy'].append(valid_logs['accuracy']) loggers['precision'].append(train_logs['precision']) loggers['val_precision'].append(valid_logs['precision']) loggers['recall'].append(train_logs['recall']) loggers['val_recall'].append(valid_logs['recall']) loggers['fscore'].append(train_logs['fscore']) loggers['val_fscore'].append(valid_logs['fscore']) # save logs to csv df = pd.DataFrame(loggers) df.to_csv(os.path.join(config.RESULT_DIR, 'logs.csv'), index=False) # plot logs plot_logs(loggers, 'loss', config.RESULT_DIR) plot_logs(loggers, 'accuracy', config.RESULT_DIR) plot_logs(loggers, 'precision', config.RESULT_DIR) plot_logs(loggers, 'recall', config.RESULT_DIR) plot_logs(loggers, 'fscore', config.RESULT_DIR) ## callbacks (save model, change lr, etc.) ## # model checkpoint if max_score < valid_logs['fscore']: max_score = valid_logs['fscore'] torch.save(model, os.path.join(config.RESULT_DIR, 'best_model.pth')) print('Model saved.') # learning rate schedule if config.SCHEDULER == 'ReduceLROnPlateau': scheduler.step(valid_logs['loss']) else: scheduler.step() # save last epoch model torch.save(model, os.path.join(config.RESULT_DIR, 'last_model.pth'))
def main(): parser = argparse.ArgumentParser( description='Train the Model on images and target masks', formatter_class=argparse.ArgumentDefaultsHelpFormatter) parser.add_argument('-g_num', '--gpu_num', metavar='G', type=str, default='0', help='GPU', dest='gpu_num') parser.add_argument('-d', '--device', metavar='G', type=str, default='cpu', help='GPU', dest='device') parser.add_argument('-train_dir', '--train_data_dir', type=str, default='/Users/ryeon/Desktop/segmentation_data/', help='dataset dir', dest='train_dir') parser.add_argument('-valid_dir', '--valid_data_dir', type=str, default='/Users/ryeon/Desktop/valid_final/result/', help='dataset dir', dest='valid_dir') parser.add_argument('-train_epoch', type=int, default=50, help='validation epoch') parser.add_argument('-valid_epoch', type=int, default=1, help='train epoch') parser.add_argument('-pretrained', type=str, default=None, help='pretrained model.pth') parser.add_argument('-encoder', type=str, default='mobilenet_v2', help='Encoder') parser.add_argument('-encoder_weight', type=str, default='imagenet', help='Encoder') parser.add_argument('-activation', type=str, default='sigmoid', help='Encoder') parser.add_argument('-simple', type=bool, default=False, help='') parser.add_argument('-trans', type=bool, default=False, help='') parser.add_argument('-color', type=bool, default=False, help='') parser.add_argument('-trans_color', type=bool, default=False, help='') args = parser.parse_args() print(args) if args.device == 'cuda': os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_num ENCODER = args.encoder ENCODER_WEIGHTS = args.encoder_weight CLASSES = ['bubble'] ACTIVATION = args.activation DEVICE = args.device CHANNEL = 3 model = smp.Unet( encoder_name=ENCODER, encoder_weights=ENCODER_WEIGHTS, in_channels=CHANNEL, classes=len(CLASSES), activation=ACTIVATION, ) loss = smp.utils.losses.DiceLoss() metrics = [ smp.utils.metrics.IoU(threshold=0.5), ] optimizer = torch.optim.Adam([ dict(params=model.parameters(), lr=0.0001), ]) DATA_DIR = args.train_dir x_train_dir = os.path.join(DATA_DIR, 'image') y_train_dir = os.path.join(DATA_DIR, 'mask') VALID_DIR = args.valid_dir preprocessing_fn = smp.encoders.get_preprocessing_fn( ENCODER, ENCODER_WEIGHTS) train_dataset = Dataset( x_train_dir, y_train_dir, augmentation=get_training_augmentation(), preprocessing=get_preprocessing(preprocessing_fn), classes=CLASSES, train=True, ) train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=0) train_epoch = smp.utils.train.TrainEpoch( model, loss=loss, metrics=metrics, optimizer=optimizer, device=DEVICE, verbose=True, ) max_score = 0 if args.pretrained != None: if args.device == 'cuda': model = torch.load(args.pretrained) else: model = torch.load(args.pretrained, map_location=args.device) for i in range(0, args.train_epoch): print(i, args.train_epoch) save_model_name = './train_model_mob_trans_' if i % args.valid_epoch == 0: print('\nValid: {}'.format(i)) bubble_mask_with_segmentation(VALID_DIR, model, i, device=args.device) print('\nEpoch: {}'.format(i)) train_logs = train_epoch.run(train_loader) # do something (save model, change lr, etc.) if max_score < train_logs['iou_score']: max_score = train_logs['iou_score'] os.makedirs('./result/model/', exist_ok=True) save_model_name = './result/model/' + save_model_name + str( i) + '.pth' torch.save(model, save_model_name) print('Model saved!') if i == 25: optimizer.param_groups[0]['lr'] = 1e-5 print('Decrease decoder learning rate to 1e-5!')
def Setup(self): ''' User function: Setup all the parameters Args: None Returns: None ''' preprocess_input = sm.get_preprocessing( self.system_dict["params"]["backbone"]) # define network parameters self.system_dict["local"]["n_classes"] = 1 if len( self.system_dict["dataset"]["train"] ["classes_to_train"]) == 1 else ( len(self.system_dict["dataset"]["train"]["classes_to_train"]) + 1) # case for binary and multiclass segmentation activation = 'sigmoid' if self.system_dict["local"][ "n_classes"] == 1 else 'softmax' #create model if (self.system_dict["params"]["model"] == "Unet"): self.system_dict["local"]["model"] = sm.Unet( self.system_dict["params"]["backbone"], classes=self.system_dict["local"]["n_classes"], activation=activation) elif (self.system_dict["params"]["model"] == "FPN"): self.system_dict["local"]["model"] = sm.FPN( self.system_dict["params"]["backbone"], classes=self.system_dict["local"]["n_classes"], activation=activation) elif (self.system_dict["params"]["model"] == "Linknet"): self.system_dict["local"]["model"] = sm.Linknet( self.system_dict["params"]["backbone"], classes=self.system_dict["local"]["n_classes"], activation=activation) elif (self.system_dict["params"]["model"] == "PSPNet"): self.system_dict["local"]["model"] = sm.PSPNet( self.system_dict["params"]["backbone"], classes=self.system_dict["local"]["n_classes"], activation=activation) # define optomizer optim = keras.optimizers.Adam(self.system_dict["params"]["lr"]) # Segmentation models losses can be combined together by '+' and scaled by integer or float factor dice_loss = sm.losses.DiceLoss() focal_loss = sm.losses.BinaryFocalLoss() if self.system_dict["local"][ "n_classes"] == 1 else sm.losses.CategoricalFocalLoss() total_loss = dice_loss + (1 * focal_loss) # actulally total_loss can be imported directly from library, above example just show you how to manipulate with losses # total_loss = sm.losses.binary_focal_dice_loss # or sm.losses.categorical_focal_dice_loss metrics = [ sm.metrics.IOUScore(threshold=0.5), sm.metrics.FScore(threshold=0.5) ] # compile keras model with defined optimozer, loss and metrics self.system_dict["local"]["model"].compile(optim, total_loss, metrics) # Dataset for train images train_dataset = Dataset( self.system_dict["dataset"]["train"]["img_dir"], self.system_dict["dataset"]["train"]["mask_dir"], self.system_dict["dataset"]["train"]["classes_dict"], classes_to_train=self.system_dict["dataset"]["train"] ["classes_to_train"], augmentation=get_training_augmentation(), preprocessing=get_preprocessing(preprocess_input), ) if (self.system_dict["params"]["image_shape"][0] % 32 != 0): self.system_dict["params"]["image_shape"][0] += ( 32 - self.system_dict["params"]["image_shape"][0] % 32) if (self.system_dict["params"]["image_shape"][1] % 32 != 0): self.system_dict["params"]["image_shape"][1] += ( 32 - self.system_dict["params"]["image_shape"][1] % 32) # Dataset for validation images if (self.system_dict["dataset"]["val"]["status"]): valid_dataset = Dataset( self.system_dict["dataset"]["val"]["img_dir"], self.system_dict["dataset"]["val"]["mask_dir"], self.system_dict["dataset"]["train"]["classes_dict"], classes_to_train=self.system_dict["dataset"]["train"] ["classes_to_train"], augmentation=get_validation_augmentation( self.system_dict["params"]["image_shape"][0], self.system_dict["params"]["image_shape"][1]), preprocessing=get_preprocessing(preprocess_input), ) else: valid_dataset = Dataset( self.system_dict["dataset"]["train"]["img_dir"], self.system_dict["dataset"]["train"]["mask_dir"], self.system_dict["dataset"]["train"]["classes_dict"], classes_to_train=self.system_dict["dataset"]["train"] ["classes_to_train"], augmentation=get_validation_augmentation( self.system_dict["params"]["image_shape"][0], self.system_dict["params"]["image_shape"][1]), preprocessing=get_preprocessing(preprocess_input), ) self.system_dict["local"]["train_dataloader"] = Dataloder( train_dataset, batch_size=self.system_dict["params"]["batch_size"], shuffle=True) self.system_dict["local"]["valid_dataloader"] = Dataloder( valid_dataset, batch_size=1, shuffle=False)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") if len(device_ids) > 1: model = DataParallel(model) min_loss = 1 early_stop_counter = 0 for epoch in range(NUM_EPOCH): dataset = CloudTrainDataset2( train_csv, data_path, kfold_path, FOLD, phase='train', transform_type=0, preprocessing=utils.get_preprocessing(preprocessing_fn), ) dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS) train_num = len(dataset) train_loss = 0 for _, imgs, masks, _ in tqdm(dataloader): imgs = imgs.float().cuda() masks = masks.float().cuda() outputs = model(imgs) loss = criterion(outputs, masks) train_loss += (loss.item() * imgs.shape[0]) optimizer.zero_grad()
train = pd.read_csv(f'train_preprocessed.csv') train_ids = pd.read_csv(f'./folds/fold_{args.fold}_train.csv').values.ravel() valid_ids = pd.read_csv(f'./folds/fold_{args.fold}_val.csv').values.ravel() num_workers = 4 bs = args.bs train_dataset = CloudDataset(df=train, image_size=(args.size, args.size * 2), path=path, datatype='train', preload=False, img_ids=train_ids, filter_bad_images=True, transforms=get_training_augmentation( size=(args.size, args.size * 2), p=0.5), preprocessing=get_preprocessing(preprocessing_fn)) valid_dataset = CloudDataset(df=train, image_size=(args.size, args.size * 2), path=path, datatype='valid', preload=False, img_ids=valid_ids, filter_bad_images=True, transforms=get_validation_augmentation( (args.size, args.size * 2)), preprocessing=get_preprocessing(preprocessing_fn)) train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=num_workers)