def main(): pyplot.style.use("bmh") base_path = Path.home() / "/Data" / "PennFudanPed" save_model_path = PROJECT_APP_PATH.user_data / 'models' / "penn_fudan_ped_seg.model" train_model = False eval_model = not train_model SEED = 87539842 batch_size = 8 num_workers = 1 # os.cpu_count() learning_rate = 0.01 torch_seed(SEED) train_set = PennFudanDataset(base_path, Split.Training) train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers) valid_loader = DataLoader( PennFudanDataset(base_path, Split.Validation), batch_size=batch_size, shuffle=False, num_workers=num_workers, ) model = SkipHourglassFission( input_channels=train_set.predictor_shape[-1], output_heads=(train_set.response_shape[-1], ), encoding_depth=1, ) model.to(global_torch_device()) if train_model: if save_model_path.exists(): model.load_state_dict(torch.load(str(save_model_path))) print("loading saved model") with TorchTrainSession(model): criterion = BCEDiceLoss(eps=1.0) optimiser = torch.optim.SGD(model.parameters(), lr=learning_rate) scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimiser, T_max=7, eta_min=learning_rate / 100, last_epoch=-1) model = train_person_segmenter( model, train_loader, valid_loader, criterion, optimiser, scheduler, save_model_path, ) if eval_model: if save_model_path.exists(): model.load_state_dict(torch.load(str(save_model_path))) print("loading saved model") with TorchDeviceSession(global_torch_device(cuda_if_available=False), model): with torch.no_grad(): with TorchCacheSession(): with TorchEvalSession(model): valid_masks = [] a = (350, 525) tr = min(len(valid_loader.dataset) * 4, 2000) probabilities = numpy.zeros((tr, *a), dtype=numpy.float32) for sample_i, (data, target) in enumerate( tqdm(valid_loader)): data = data.to(global_torch_device()) target = target.cpu().detach().numpy() outpu, *_ = model(data) outpu = torch.sigmoid(outpu).cpu().detach().numpy() for p in range(data.shape[0]): output, mask = outpu[p], target[p] """ for m in mask: valid_masks.append(cv2_resize(m, a)) for probability in output: probabilities[sample_i, :, :] = cv2_resize(probability, a) sample_i += 1 """ if sample_i >= tr - 1: break if sample_i >= tr - 1: break f, ax = pyplot.subplots(3, 3, figsize=(24, 12)) for i in range(3): ax[0, i].imshow(valid_masks[i], vmin=0, vmax=1) ax[0, i].set_title("Original", fontsize=14) ax[1, i].imshow(valid_masks[i], vmin=0, vmax=1) ax[1, i].set_title("Target", fontsize=14) ax[2, i].imshow(probabilities[i], vmin=0, vmax=1) ax[2, i].set_title("Prediction", fontsize=14) pyplot.show()
def train_mnist(load_earlier=False, train=True, denoise: bool = True): """ :param load_earlier: :type load_earlier: :param train: :type train:""" seed = 251645 batch_size = 32 tqdm.monitor_interval = 0 learning_rate = 3e-3 lr_sch_step_size = int(10e4 // batch_size) lr_sch_gamma = 0.1 unet_depth = 3 unet_start_channels = 16 input_channels = 1 output_channels = (input_channels,) home_path = PROJECT_APP_PATH model_file_ending = ".model" model_base_path = ensure_existence(PROJECT_APP_PATH.user_data / "unet_mnist") interrupted_name = "INTERRUPTED_BEST" interrupted_path = model_base_path / f"{interrupted_name}{model_file_ending}" torch.manual_seed(seed) device = global_torch_device() img_transform = transforms.Compose( [ transforms.ToTensor(), MinMaxNorm(), transforms.Lambda(lambda tensor: torch.round(tensor)), # transforms.RandomErasing() ] ) dataset = MNIST( PROJECT_APP_PATH.user_data / "mnist", transform=img_transform, download=True ) data_iter = iter( cycle(DataLoader(dataset, batch_size=batch_size, shuffle=True, pin_memory=True)) ) data_iter = to_device_iterator(data_iter, device) model = SkipHourglassFission( input_channels=input_channels, output_heads=output_channels, encoding_depth=unet_depth, start_channels=unet_start_channels, ).to(global_torch_device()) optimiser_ft = optim.Adam(model.parameters(), lr=learning_rate) exp_lr_scheduler = optim.lr_scheduler.StepLR( optimiser_ft, step_size=lr_sch_step_size, gamma=lr_sch_gamma ) if load_earlier: _list_of_files = list( model_base_path.rglob(f"{interrupted_name}{model_file_ending}") ) if not len(_list_of_files): print( f"found no trained models under {model_base_path}{os.path.sep}**{os.path.sep}{interrupted_name}{model_file_ending}" ) exit(1) latest_model_path = str(max(_list_of_files, key=os.path.getctime)) print(f"loading previous model: {latest_model_path}") if latest_model_path is not None: model.load_state_dict(torch.load(latest_model_path)) if train: with TensorBoardPytorchWriter(home_path.user_log / str(time.time())) as writer: model = training( model, data_iter, optimiser_ft, exp_lr_scheduler, writer, interrupted_path, denoise=denoise, ) torch.save( model.state_dict(), model_base_path / f"unet_mnist_final{model_file_ending}", ) else: inference(model, data_iter, denoise=denoise) torch.cuda.empty_cache()
def main( base_path: Path = Path.home() / "Data" / "Datasets" / "PennFudanPed", train_model: bool = True, load_prev_model: bool = True, writer: Writer = TensorBoardPytorchWriter(PROJECT_APP_PATH.user_log / "instanced_person_segmentation" / f"{time.time()}"), ): """ """ # base_path = Path("/") / "encrypted_disk" / "heider" / "Data" / "PennFudanPed" base_path: Path = Path.home() / "Data3" / "PennFudanPed" # base_path = Path('/media/heider/OS/Users/Christian/Data/Datasets/') / "PennFudanPed" pyplot.style.use("bmh") save_model_path = ( ensure_existence(PROJECT_APP_PATH.user_data / "models") / "instanced_penn_fudan_ped_seg.model") eval_model = not train_model SEED = 9221 batch_size = 32 num_workers = 0 encoding_depth = 2 learning_rate = 6e-6 # sequence 6e-2 6e-3 6e-4 6e-5 seed_stack(SEED) train_set = PennFudanDataset( base_path, SplitEnum.training, return_variant=PennFudanDataset.PennFudanReturnVariantEnum.instanced, ) train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers) valid_loader = DataLoader( PennFudanDataset( base_path, SplitEnum.validation, return_variant=PennFudanDataset.PennFudanReturnVariantEnum. instanced, ), batch_size=batch_size, shuffle=False, num_workers=num_workers, ) model = SkipHourglassFission( input_channels=train_set.predictor_shape[-1], output_heads=(train_set.response_shape[-1], ), encoding_depth=encoding_depth, ) model.to(global_torch_device()) if load_prev_model and save_model_path.exists(): model.load_state_dict(torch.load(str(save_model_path))) print("loading saved model") if train_model: with TorchTrainSession(model): criterion = BCEDiceLoss() # optimiser = torch.optim.SGD(model.parameters(), lr=learning_rate) optimiser = torch.optim.Adam(model.parameters(), lr=learning_rate) # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimiser, T_max=7, eta_min=learning_rate / 100, last_epoch=-1 ) model = train_person_segmentor( model, train_loader, valid_loader, criterion, optimiser, save_model_path=save_model_path, learning_rate=learning_rate, writer=writer, ) if eval_model: validate_model(model, valid_loader)
def main(): args = argparse.ArgumentParser() args.add_argument("-i", action="store_false") options = args.parse_args() seed = 42 batch_size = 8 # 12 depth = 4 # 5 segmentation_channels = 3 tqdm.monitor_interval = 0 learning_rate = 3e-3 lr_sch_step_size = int(1000 // batch_size) + 4 lr_sch_gamma = 0.1 model_start_channels = 16 home_path = Path.home() / "Models" / "Vision" base_path = home_path / str(time.time()) best_model_path = "INTERRUPTED_BEST.pth" interrupted_path = str(base_path / best_model_path) writer = TensorBoardPytorchWriter(str(base_path)) env = CameraObservationWrapper() torch.manual_seed(seed) env.seed(seed) device = global_torch_device() aeu_model = SkipHourglassFission( segmentation_channels, (segmentation_channels,), encoding_depth=depth, start_channels=model_start_channels, ) aeu_model = aeu_model.to(global_torch_device()) optimizer_ft = optim.Adam(aeu_model.parameters(), lr=learning_rate) exp_lr_scheduler = lr_scheduler.StepLR( optimizer_ft, step_size=lr_sch_step_size, gamma=lr_sch_gamma ) data_iter = iter(neodroid_camera_data_iterator(env, device, batch_size)) if options.i: trained_aeu_model = train_model( aeu_model, data_iter, optimizer_ft, exp_lr_scheduler, writer, interrupted_path, ) test_model(trained_aeu_model, data_iter) else: _list_of_files = list(home_path.glob("*")) latest_model_path = ( str(max(_list_of_files, key=os.path.getctime)) + f"/{best_model_path}" ) print("loading previous model: " + latest_model_path) test_model(aeu_model, data_iter, load_path=latest_model_path) torch.cuda.empty_cache() env.close() writer.close()
def export_detection_model( model_export_path: Path = ensure_existence( PROJECT_APP_PATH.user_data / "penn_fudan_segmentation" ) / "seg_skip_fis", SEED: int = 87539842, ) -> None: """ :param model_export_path: :type model_export_path: :return: :rtype:""" model = OutputActivationModule( SkipHourglassFission(input_channels=3, output_heads=(1,), encoding_depth=1) ) with TorchDeviceSession(device=global_torch_device("cpu"), model=model): with TorchEvalSession(model): seed_stack(SEED) # standard PyTorch mean-std input image normalization transform = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) frame_g = frame_generator(cv2.VideoCapture(0)) for image in tqdm(frame_g): example_input = ( transform(image).unsqueeze(0).to(global_torch_device()), ) try: traced_script_module = torch.jit.trace( model, example_input, # strict=strict_jit, check_inputs=( transform(next(frame_g)) .unsqueeze(0) .to(global_torch_device()), transform(next(frame_g)) .unsqueeze(0) .to(global_torch_device()), ), ) exp_path = model_export_path.with_suffix(".traced") traced_script_module.save(str(exp_path)) print( f"Traced Ops used {torch.jit.export_opnames(traced_script_module)}" ) sprint( f"Successfully exported JIT Traced model at {exp_path}", color="green", ) except Exception as e_i: sprint(f"Torch JIT Trace export does not work!, {e_i}", color="red") break