def test_nested_model_sessions(): model = torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.Dropout(0.1)) print(model.training) with TorchEvalSession(model): print(model.training) with TorchTrainSession(model): print(model.training) with TorchEvalSession(model): print(model.training) with TorchTrainSession(model): print(model.training) with TorchEvalSession(model): print(model.training) print(model.training)
def run_seg_traced_webcam_demo(): """ :return: :rtype:""" import torch import io load_path = (PROJECT_APP_PATH.user_data / "penn_fudan_segmentation" / "seg_skip_fis").with_suffix(".traced") # print(load_path) # torch.jit.load(str(load_path)) with open(str(load_path), "rb") as f: # Load ScriptModule from io.BytesIO object buffer = io.BytesIO(f.read()) model = torch.jit.load(buffer) # Load all tensors to the original device transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) from matplotlib.pyplot import imshow with TorchDeviceSession(device=global_torch_device("cpu"), model=model): with TorchEvalSession(model): for image in tqdm(frame_generator(cv2.VideoCapture(0))): result = model( transform(image).unsqueeze(0).to(global_torch_device()))[0] imshow(result[0][0].numpy(), vmin=0.0, vmax=1.0) show()
def maskrcnn_evaluate( model: Module, data_loader: DataLoader, *, device=global_torch_device(), writer: Writer = None, ) -> CocoEvaluator: """ Args: model: data_loader: device: writer: Returns: """ n_threads = torch.get_num_threads() # FIXME remove this and make paste_masks_in_image run on the GPU torch.set_num_threads(1) cpu_device = torch.device("cpu") coco_evaluator = CocoEvaluator( get_coco_api_from_dataset(data_loader.dataset), get_iou_types(model)) with torch.no_grad(): with TorchEvalSession(model): for image, targets in tqdm.tqdm(data_loader): image = [img.to(device) for img in image] targets = [{k: v.to(device) for k, v in t.items()} for t in targets] torch.cuda.synchronize(device) model_time = time.time() outputs = model(image) outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs] model_time = time.time() - model_time res = { target["image_id"].item(): output for target, output in zip(targets, outputs) } evaluator_time = time.time() coco_evaluator.update(res) evaluator_time = time.time() - evaluator_time if writer: writer.scalar("model_time", model_time) writer.scalar("evaluator_time", evaluator_time) coco_evaluator.synchronize_between_processes() coco_evaluator.accumulate() coco_evaluator.summarize() torch.set_num_threads(n_threads) return coco_evaluator
def main(model_name: str = "maskrcnn_pennfudanped", score_threshold=0.55): base_path = PROJECT_APP_PATH.user_data / 'maskrcnn' dataset_root = Path.home() / "Data" torch_seed(3825) dataset = PennFudanDataset(dataset_root / "PennFudanPed", Split.Training) categories = dataset.categories if True: model = load_model(model_name=model_name, model_directory=base_path / 'models') else: model = get_pretrained_instance_segmentation_maskrcnn( dataset.response_channels) model.to(global_torch_device()) cpu_device = torch.device("cpu") with torch.no_grad(): with TorchEvalSession(model): for image in tqdm( to_tensor_generator( frame_generator(cv2.VideoCapture(0)), device=global_torch_device(), )): prediction = model( # torch_vision_normalize_batch_nchw( uint_hwc_to_chw_float_tensor(image).unsqueeze(0) # ) )[0] (boxes, labels, scores) = ( prediction["boxes"].to(cpu_device).numpy(), prediction["labels"].to(cpu_device).numpy(), torch.sigmoid( prediction["scores"]).to(cpu_device).numpy(), ) indices = scores > score_threshold cv2.namedWindow(model_name, cv2.WINDOW_NORMAL) cv2.imshow( model_name, draw_bounding_boxes( quick_to_pil_image(image), boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, )) if cv2.waitKey(1) == 27: break # esc to quit
def test_model( model: VAE, epoch_i: int, metric_writer: Writer, loader: DataLoader, save_images: bool = True, ): global LOWEST_L with TorchEvalSession(model): test_accum_loss = 0 with torch.no_grad(): for i, (original, labels, *_) in enumerate(loader): original = original.to(global_torch_device()) reconstruction, mean, log_var = model(original) loss = loss_function(reconstruction, original, mean, log_var).item() test_accum_loss += loss metric_writer.scalar("test_loss", test_accum_loss) if save_images: if i == 0: n = min(original.size(0), 8) comparison = torch.cat( [original[:n], reconstruction[:n]]) save_image( comparison.cpu(), # Torch save images str(BASE_PATH / f"reconstruction_{str(epoch_i)}.png"), nrow=n, ) """ scatter_plot_encoding_space(str(BASE_PATH / f'encoding_space_{str(epoch_i)}.png'), mean.to('cpu').numpy(), log_var.to('cpu').numpy(), labels) """ break # test_loss /= len(loader.dataset) test_accum_loss /= loader.batch_size print(f"====> Test set loss: {test_accum_loss:.4f}") torch.save(model.state_dict(), BASE_PATH / f"model_state_dict{str(epoch_i)}.pth") if LOWEST_L > test_accum_loss: LOWEST_L = test_accum_loss torch.save(model.state_dict(), BASE_PATH / f"best_state_dict.pth")
def run_seg_traced_webcam_demo(): """ :param categories: :type categories: :param cfg: :type cfg: :param model_ckpt: :type model_ckpt: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype: """ import torch import io torch.jit.load("seg_skip_fis.traced") with open( "seg_skip_fis.traced", "rb" ) as f: # Load ScriptModule from io.BytesIO object buffer = io.BytesIO(f.read()) model = torch.jit.load(buffer) # Load all tensors to the original device transform = transforms.Compose( [ transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) with TorchDeviceSession( device=global_torch_device(cuda_if_available=False), model=model ): with TorchEvalSession(model): for image in tqdm(frame_generator(cv2.VideoCapture(0))): result = model(transform(image).unsqueeze(0).to(global_torch_device()))[ 0 ] print(result) from matplotlib.pyplot import imshow imshow(result[0][0].numpy(), vmin=0.0, vmax=1.0) show()
def validate_model(model, valid_loader): with TorchDeviceSession(global_torch_device("cpu"), model): with torch.no_grad(): with TorchCacheSession(): with TorchEvalSession(model): valid_masks = [] out_data = [] a = (256, 256) tr = min(len(valid_loader.dataset) * 4, 2000) probabilities = [] for sample_i, (data, target) in enumerate(tqdm(valid_loader)): data = data.to(global_torch_device()) outpu, *_ = model(data) for m, d, p in zip( target.cpu().detach().numpy(), data.cpu().detach().numpy(), torch.sigmoid(outpu).cpu().detach().numpy(), ): out_data.append(cv2_resize(chw_to_hwc(d), a)) valid_masks.append(cv2_resize(m[0], a)) probabilities.append(cv2_resize(p[0], a)) sample_i += 1 if sample_i >= tr - 1: break if sample_i >= tr - 1: break min_a = min(3, len(out_data)) f, ax = pyplot.subplots(min_a, 3, figsize=(24, 12)) # assert len(valid_masks)>2, f'{len(valid_masks), tr}' for i in range(min_a): ax[0, i].imshow(out_data[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 single_epoch_evaluation( model: Module, evaluation_loader: DataLoader, subset: Split, *, epoch: int = None, writer: Writer = None, device: torch.device = global_torch_device()) -> float: correct = 0 num_batches = len(evaluation_loader) with TorchEvalSession(model): for data, target in tqdm(evaluation_loader, desc=f'{subset} batch #', total=num_batches): correct += model(data.to(device)).argmax(dim=-1).squeeze().eq( target.to(device)).sum().item() acc = correct / len(evaluation_loader.dataset) if writer: writer.scalar(f'{subset}_accuracy', acc, epoch) return acc
def inference( model: Module, data_iterator: Iterator, denoise: bool = True, num: int = 3 ) -> None: """ :param model: :type model: :param data_iterator: :type data_iterator:""" with torch.no_grad(): with TorchEvalSession(model): img, target = next(data_iterator) if denoise: model_input = img + torch.normal( mean=0.0, std=0.1, size=img.shape, device=global_torch_device() ) else: model_input = img model_input = torch.clamp(model_input, 0.0, 1.0) pred, *_ = model(model_input) plot_side_by_side( [ pred.squeeze(1)[:num].cpu().numpy(), model_input.squeeze(1)[:num].cpu().numpy(), ] ) pyplot.show() return for i, (s, j, label) in enumerate( zip(pred.cpu().numpy(), model_input.cpu().numpy(), target) ): pyplot.imshow(j[0]) pyplot.title(f"sample_{i}, category: {label}") pyplot.show() # plot_side_by_side(s) pyplot.imshow(s[0]) pyplot.title(f"sample_{i}, category: {label}") pyplot.show() break
def train_siamese( model, optimiser, criterion, *, writer: Writer = MockWriter(), train_number_epochs, data_dir, train_batch_size, model_name, save_path, save_best=False, img_size, validation_interval: int = 1, ): """ :param data_dir: :type data_dir: :param optimiser: :type optimiser: :param criterion: :type criterion: :param writer: :type writer: :param model_name: :type model_name: :param save_path: :type save_path: :param save_best: :type save_best: :param model: :type model: :param train_number_epochs: :type train_number_epochs: :param train_batch_size: :type train_batch_size: :return: :rtype: Parameters ---------- img_size validation_interval""" train_dataloader = DataLoader( TripletDataset( data_path=data_dir, transform=transforms.Compose([ transforms.Grayscale(), transforms.Resize(img_size), transforms.ToTensor(), ]), split=SplitEnum.training, ), shuffle=True, num_workers=0, batch_size=train_batch_size, ) valid_dataloader = DataLoader( TripletDataset( data_path=data_dir, transform=transforms.Compose([ transforms.Grayscale(), transforms.Resize(img_size), transforms.ToTensor(), ]), split=SplitEnum.validation, ), shuffle=True, num_workers=0, batch_size=train_batch_size, ) best = math.inf E = tqdm(range(0, train_number_epochs)) batch_counter = count() for epoch in E: for tss in train_dataloader: batch_i = next(batch_counter) with TorchTrainSession(model): optimiser.zero_grad() loss_contrastive = criterion(*model( *[t.to(global_torch_device()) for t in tss])) loss_contrastive.backward() optimiser.step() a = loss_contrastive.cpu().item() writer.scalar("train_loss", a, batch_i) if batch_counter.__next__() % validation_interval == 0: with TorchEvalSession(model): for tsv in valid_dataloader: o = model(*[t.to(global_torch_device()) for t in tsv]) a_v = criterion(*o).cpu().item() valid_positive_acc = (accuracy( distances=pairwise_distance(o[0], o[1]), is_diff=0).cpu().item()) valid_negative_acc = (accuracy( distances=pairwise_distance(o[0], o[2]), is_diff=1).cpu().item()) valid_acc = numpy.mean( (valid_negative_acc, valid_positive_acc)) writer.scalar("valid_loss", a_v, batch_i) writer.scalar("valid_positive_acc", valid_positive_acc, batch_i) writer.scalar("valid_negative_acc", valid_negative_acc, batch_i) writer.scalar("valid_acc", valid_acc, batch_i) if a_v < best: best = a_v print(f"new best {best}") if save_best: save_model_parameters( model, optimiser=optimiser, model_name=model_name, save_directory=save_path, ) E.set_description( f"Epoch number {epoch}, Current train loss {a}, valid loss {a_v}, valid acc {valid_acc}" ) return model
def pred_target_train_model( model, train_iterator, criterion, optimizer, scheduler, writer, interrupted_path, test_data_iterator=None, num_updates: int = 250000, early_stop=None, ) -> torch.nn.Module: best_model_wts = copy.deepcopy(model.state_dict()) best_val_loss = 1e10 since = time.time() try: sess = tqdm.tqdm(range(num_updates), leave=False, disable=False) val_loss = 0 update_loss = 0 val_acc = 0 last_val = None last_out = None with torch.autograd.detect_anomaly(): for update_i in sess: for phase in [Split.Training, Split.Validation]: if phase == Split.Training: with TorchTrainSession(model): input, true_label = zip(*next(train_iterator)) rgb_imgs = torch_vision_normalize_batch_nchw( uint_nhwc_to_nchw_float_batch( rgb_drop_alpha_batch_nhwc( to_tensor(input)))) true_label = to_tensor(true_label, dtype=torch.long) optimizer.zero_grad() pred = model(rgb_imgs) loss = criterion(pred, true_label) loss.backward() optimizer.step() if last_out is None: last_out = pred else: if not torch.dist(last_out, pred) > 0: print(f"Same output{last_out},{pred}") last_out = pred update_loss = loss.data.cpu().numpy() writer.scalar(f"loss/train", update_loss, update_i) if scheduler: scheduler.step() elif test_data_iterator: with TorchEvalSession(model): test_rgb_imgs, test_true_label = zip( *next(train_iterator)) test_rgb_imgs = torch_vision_normalize_batch_nchw( uint_nhwc_to_nchw_float_batch( rgb_drop_alpha_batch_nhwc( to_tensor(test_rgb_imgs)))) test_true_label = to_tensor(test_true_label, dtype=torch.long) with torch.no_grad(): val_pred = model(test_rgb_imgs) val_loss = criterion(val_pred, test_true_label) _, cat = torch.max(val_pred, -1) val_acc = torch.sum( cat == test_true_label) / float(cat.size(0)) writer.scalar(f"loss/acc", val_acc, update_i) writer.scalar(f"loss/val", val_loss, update_i) if last_val is None: last_val = cat else: if all(last_val == cat): print(f"Same val{last_val},{cat}") last_val = cat if val_loss < best_val_loss: best_val_loss = val_loss best_model_wts = copy.deepcopy( model.state_dict()) sess.write( f"New best validation model at update {update_i} with test_loss {best_val_loss}" ) torch.save(model.state_dict(), interrupted_path) if early_stop is not None and val_pred < early_stop: break sess.set_description_str(f"Update {update_i} - {phase} " f"update_loss:{update_loss:2f} " f"test_loss:{val_loss}" f"val_acc:{val_acc}") except KeyboardInterrupt: print("Interrupt") finally: pass model.load_state_dict(best_model_wts) # load best model weights time_elapsed = time.time() - since print(f"{time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s") print(f"Best val loss: {best_val_loss:3f}") return model
def run_webcam_demo( cfg: NOD, input_cfg: NOD, categories: List, model_ckpt: Path, score_threshold: float = 0.7, window_name: str = "SSD", ): """ :param categories: :type categories: :param cfg: :type cfg: :param model_ckpt: :type model_ckpt: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype: """ cpu_device = torch.device("cpu") transforms = SSDTransform(input_cfg.image_size, input_cfg.pixel_mean, split=Split.Testing) model = SingleShotDectectionNms(cfg) checkpointer = CheckPointer(model, save_dir=ensure_existence( PROJECT_APP_PATH.user_data / "results")) checkpointer.load(model_ckpt, use_latest=model_ckpt is None) print( f"Loaded weights from {model_ckpt if model_ckpt else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) with TorchEvalSession(model): for image in tqdm(frame_generator(cv2.VideoCapture(0))): result = model( transforms(image)[0].unsqueeze(0).to(global_torch_device())) height, width, *_ = image.shape result.boxes[:, 0::2] *= width / result.img_width.cpu().item() result.boxes[:, 1::2] *= height / result.img_height.cpu().item() (boxes, labels, scores) = ( result.boxes.to(cpu_device).numpy(), result.labels.to(cpu_device).numpy(), result.scores.to(cpu_device).numpy(), ) indices = scores > score_threshold cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.imshow( window_name, draw_bounding_boxes( image, boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, score_font=ImageFont.truetype( PACKAGE_DATA_PATH / "Lato-Regular.ttf", 24, ), ).astype(numpy.uint8), ) if cv2.waitKey(1) == 27: break # esc to quit
def predictor_response_train_model( model, *, train_iterator, criterion, optimizer, scheduler, writer, interrupted_path, val_data_iterator=None, num_updates: int = 250000, device=global_torch_device(), early_stop=None, ): """ :param model: :param train_iterator: :param criterion: :param optimizer: :param scheduler: :param writer: :param interrupted_path: :param val_data_iterator: :param num_updates: :param device: :param early_stop: :return: """ best_model_wts = copy.deepcopy(model.state_dict()) best_val_loss = 1e10 since = time.time() try: sess = tqdm(range(num_updates), leave=False, disable=False) val_loss = 0 update_loss = 0 val_acc = 0 last_val = None last_out = None with torch.autograd.detect_anomaly(): for update_i in sess: for phase in [Split.Training, Split.Validation]: if phase == Split.Training: with TorchTrainSession(model): input, true_label = next(train_iterator) rgb_imgs = to_tensor( input, dtype=torch.float, device=device ).repeat(1, 3, 1, 1) true_label = to_tensor( true_label, dtype=torch.long, device=device ) optimizer.zero_grad() pred = model(rgb_imgs) loss = criterion(pred, true_label) loss.backward() optimizer.step() update_loss = loss.data.cpu().numpy() writer.scalar(f"loss/train", update_loss, update_i) if scheduler: scheduler.step() elif val_data_iterator: with TorchEvalSession(model): test_rgb_imgs, test_true_label = next(val_data_iterator) test_rgb_imgs = to_tensor( test_rgb_imgs, dtype=torch.float, device=device ).repeat(1, 3, 1, 1) test_true_label = to_tensor( test_true_label, dtype=torch.long, device=device ) with torch.no_grad(): val_pred = model(test_rgb_imgs) val_loss = criterion(val_pred, test_true_label) _, cat = torch.max(val_pred, -1) val_acc = torch.sum(cat == test_true_label) / float( cat.size(0) ) writer.scalar(f"loss/acc", val_acc, update_i) writer.scalar(f"loss/val", val_loss, update_i) if val_loss < best_val_loss: best_val_loss = val_loss best_model_wts = copy.deepcopy(model.state_dict()) sess.write( f"New best validation model at update {update_i} with best_val_loss {best_val_loss}" ) torch.save(model.state_dict(), interrupted_path) if early_stop is not None and val_pred < early_stop: break sess.set_description_str( f"Update {update_i} - {phase} " f"update_loss:{update_loss:2f} " f"val_loss:{val_loss}" f"val_acc:{val_acc}" ) except KeyboardInterrupt: print("Interrupt") finally: pass model.load_state_dict(best_model_wts) # load best model weights time_elapsed = time.time() - since print(f"{time_elapsed // 60:.0f}m {time_elapsed % 60:.0f}s") print(f"Best val loss: {best_val_loss:3f}") return model
def train_person_segmentor( model: torch.nn.Module, train_loader: torch.utils.data.DataLoader, valid_loader: torch.utils.data.DataLoader, criterion: callable, optimiser: torch.optim.Optimizer, *, save_model_path: Path, learning_rate: Number = 6e-2, scheduler: torch.optim.lr_scheduler = None, n_epochs: int = 100, writer: ImageWriterMixin = MockWriter(), ): """ :param model: :type model: :param train_loader: :type train_loader: :param valid_loader: :type valid_loader: :param criterion: :type criterion: :param optimiser: :type optimiser: :param scheduler: :type scheduler: :param save_model_path: :type save_model_path: :param n_epochs: :type n_epochs: :return: :rtype:""" valid_loss_min = numpy.Inf # track change in validation loss assert n_epochs > 0, n_epochs E = tqdm(range(1, n_epochs + 1)) for epoch_i in E: train_loss = 0.0 valid_loss = 0.0 with TorchTrainSession(model): for data, target in tqdm(train_loader): output, *_ = model(data.to(global_torch_device())) loss = criterion(output, target.to(global_torch_device()).float()) optimiser.zero_grad() loss.backward() optimiser.step() train_loss += loss.cpu().item() * data.size(0) with TorchEvalSession(model): with torch.no_grad(): for data, target in tqdm(valid_loader): target = target.float() ( output, *_, ) = model( # forward pass: compute predicted outputs by passing inputs to the model data.to(global_torch_device())) validation_loss = criterion( output, target.to( global_torch_device())) # calculate the batch loss writer.scalar( "dice_validation", dice_loss(output, target.to(global_torch_device())), ) valid_loss += validation_loss.detach().cpu().item( ) * data.size(0) # update average validation loss writer.image("input", data, epoch_i) # write the last batch writer.image("truth", target, epoch_i) # write the last batch writer.image("prediction", torch.sigmoid(output), epoch_i) # write the last batch # calculate average losses train_loss = train_loss / len(train_loader.dataset) valid_loss = valid_loss / len(valid_loader.dataset) # save model if validation loss has decreased if valid_loss <= valid_loss_min: print( f"Validation loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f}). Saving model ..." ) torch.save(model.state_dict(), save_model_path) valid_loss_min = valid_loss if scheduler: scheduler.step() optimiser, scheduler = reschedule_learning_rate( model, optimiser, epoch_i, scheduler, starting_learning_rate=learning_rate, ) # print training/validation statistics current_lr = next(iter(optimiser.param_groups))["lr"] E.set_description(f"Epoch: {epoch_i} " f"Training Loss: {train_loss:.6f} " f"Validation Loss: {valid_loss:.6f} " f"Learning rate: {current_lr:.6f}") writer.scalar("training_loss", train_loss) writer.scalar("validation_loss", valid_loss) writer.scalar("learning_rate", current_lr) return model
def inner_train_ssd(*, data_root: Path, cfg: NOD, model: Module, data_loader: DataLoader, optimiser: Optimizer, scheduler: WarmupMultiStepLR, check_pointer: callable, device: callable, arguments: callable, kws: NOD, ) -> Module: """ :param data_root: :type data_root: :param cfg: :type cfg: :param model: :type model: :param data_loader: :type data_loader: :param optimiser: :type optimiser: :param scheduler: :type scheduler: :param check_pointer: :type check_pointer: :param device: :type device: :param arguments: :type arguments: :param kws: :type kws: :return: :rtype: """ logger = logging.getLogger("SSD.trainer") logger.info("Start training ...") meters = MetricLogger() with TorchTrainSession(model): save_to_disk = global_distribution_rank() == 0 if kws.use_tensorboard and save_to_disk: import tensorboardX writer = tensorboardX.SummaryWriter( log_dir=str(PROJECT_APP_PATH.user_data / "results" / "tf_logs") ) else: writer = None max_iter = len(data_loader) start_iter = arguments["iteration"] start_training_time = time.time() end = time.time() for iteration, (images, targets, _) in enumerate(data_loader, start_iter): arguments["iteration"] = iteration images = images.to(device) targets = targets.to(device) loss_instance = MultiBoxLoss(neg_pos_ratio=cfg.model.neg_pos_ratio) cls_logits, bbox_pred = model(images) reg_loss, cls_loss = loss_instance( cls_logits, bbox_pred, targets.labels, targets.boxes ) loss_dict = dict(reg_loss=reg_loss, cls_loss=cls_loss) loss = sum(loss for loss in loss_dict.values()) loss_dict_reduced = reduce_loss_dict( loss_dict ) # reduce losses over all GPUs for logging purposes losses_reduced = sum(loss for loss in loss_dict_reduced.values()) meters.update(total_loss=losses_reduced, **loss_dict_reduced) optimiser.zero_grad() loss.backward() optimiser.step() scheduler.step() batch_time = time.time() - end end = time.time() meters.update(time=batch_time) if iteration % kws.log_step == 0: eta_seconds = meters.time.global_avg * (max_iter - iteration) eta_string = str(datetime.timedelta(seconds=int(eta_seconds))) logger.info( meters.delimiter.join( [ f"iter: {iteration:06d}", f"lr: {optimiser.param_groups[0]['lr']:.5f}", f"{str(meters)}", f"eta: {eta_string}", f"mem: {round(torch.cuda.max_memory_allocated() / 1024.0 / 1024.0)}M", ] ) ) if writer: global_step = iteration writer.add_scalar( "losses/total_loss", losses_reduced, global_step=global_step ) for loss_name, loss_item in loss_dict_reduced.items(): writer.add_scalar( f"losses/{loss_name}", loss_item, global_step=global_step ) writer.add_scalar( "lr", optimiser.param_groups[0]["lr"], global_step=global_step ) if iteration % kws.save_step == 0: check_pointer.save(f"model_{iteration:06d}", **arguments) if ( kws.eval_step > 0 and iteration % kws.eval_step == 0 and not iteration == max_iter ): with TorchEvalSession(model): eval_results = do_ssd_evaluation( data_root, cfg, model, distributed=kws.distributed, iteration=iteration, ) if global_distribution_rank() == 0 and writer: for eval_result, dataset in zip( eval_results, cfg.datasets.test ): write_metrics_recursive( eval_result["metrics"], "metrics/" + dataset, writer, iteration, ) check_pointer.save("model_final", **arguments) total_training_time = int( time.time() - start_training_time ) # compute training time logger.info( f"Total training time: {datetime.timedelta(seconds=total_training_time)} (" f"{total_training_time / max_iter:.4f} s / it)" ) return model
def main(): """ """ args = argparse.ArgumentParser() args.add_argument("--inference", "-i", action="store_true") args.add_argument("--continue_training", "-c", action="store_true") args.add_argument("--no_cuda", "-k", action="store_false") args.add_argument("--export", "-e", action="store_true") options = args.parse_args() timeas = str(time.time()) this_model_path = PROJECT_APP_PATH.user_data / timeas this_log = PROJECT_APP_PATH.user_log / timeas ensure_directory_exist(this_model_path) ensure_directory_exist(this_log) best_model_name = "best_validation_model.model" interrupted_path = str(this_model_path / best_model_name) torch.manual_seed(seed) if not options.no_cuda: global_torch_device("cpu") with MixedObservationWrapper() as env: env.seed(seed) train_iter = batch_generator(iter(env), batch_size) num_categories = env.sensor("Class").space.discrete_steps test_iter = train_iter model, params_to_update = squeezenet_retrain(num_categories) print(params_to_update) model = model.to(global_torch_device()) if options.continue_training: _list_of_files = list(PROJECT_APP_PATH.user_data.rglob("*.model")) 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)) criterion = torch.nn.CrossEntropyLoss().to(global_torch_device()) optimiser_ft = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum, weight_decay=wd) exp_lr_scheduler = torch.optim.lr_scheduler.StepLR(optimiser_ft, step_size=7, gamma=0.1) with TensorBoardPytorchWriter(this_log) as writer: if True: model = pred_target_train_model( model, train_iter, criterion, optimiser_ft, exp_lr_scheduler, writer, interrupted_path, test_data_iterator=test_iter, num_updates=num_updates, ) inputs, true_label = zip(*next(train_iter)) rgb_imgs = torch_vision_normalize_batch_nchw( uint_nhwc_to_nchw_float_batch( rgb_drop_alpha_batch_nhwc(to_tensor(inputs)))) predicted = torch.argmax(model(rgb_imgs), -1) true_label = to_tensor(true_label, dtype=torch.long) print(predicted, true_label) horizontal_imshow( inputs, [ f"p:{int(p)},t:{int(t)}" for p, t in zip(predicted, true_label) ], ) pyplot.show() torch.cuda.empty_cache() if options.export: with TorchEvalSession(model): example = torch.rand(1, 3, 256, 256) traced_script_module = torch.jit.trace(model.to("cpu"), example) traced_script_module.save("resnet18_v.model")
def train_person_segmenter( model, train_loader, valid_loader, criterion, optimizer, scheduler, save_model_path: Path, n_epochs: int = 100, ): """ :param model: :type model: :param train_loader: :type train_loader: :param valid_loader: :type valid_loader: :param criterion: :type criterion: :param optimizer: :type optimizer: :param scheduler: :type scheduler: :param save_model_path: :type save_model_path: :param n_epochs: :type n_epochs: :return: :rtype: """ valid_loss_min = numpy.Inf # track change in validation loss assert n_epochs > 0, n_epochs E = tqdm(range(1, n_epochs + 1)) for epoch in E: train_loss = 0.0 valid_loss = 0.0 dice_score = 0.0 with TorchTrainSession(model): for data, target in tqdm(train_loader): data, target = ( data.to(global_torch_device()), target.to(global_torch_device()), ) optimizer.zero_grad() output, *_ = model(data) output = torch.sigmoid(output) loss = criterion(output, target) loss.backward() optimizer.step() train_loss += loss.item() * data.size(0) with TorchEvalSession(model): with torch.no_grad(): for data, target in tqdm(valid_loader): data, target = ( data.to(global_torch_device()), target.to(global_torch_device()), ) output, *_ = model( data ) # forward pass: compute predicted outputs by passing inputs to the model output = torch.sigmoid(output) loss = criterion(output, target) # calculate the batch loss valid_loss += loss.item() * data.size( 0) # update average validation loss dice_cof = intersection_over_union( output.cpu().detach().numpy(), target.cpu().detach().numpy()) dice_score += dice_cof * data.size(0) # calculate average losses train_loss = train_loss / len(train_loader.dataset) valid_loss = valid_loss / len(valid_loader.dataset) dice_score = dice_score / len(valid_loader.dataset) # print training/validation statistics E.set_description(f"Epoch: {epoch}" f" Training Loss: {train_loss:.6f} " f"Validation Loss: {valid_loss:.6f} " f"Dice Score: {dice_score:.6f}") # save model if validation loss has decreased if valid_loss <= valid_loss_min: print( f"Validation loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f}). Saving model ..." ) torch.save(model.state_dict(), save_model_path) valid_loss_min = valid_loss scheduler.step() model, scheduler = reschedule_learning_rate(model, epoch, scheduler) return model
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_d( model, train_loader, valid_loader, criterion, optimiser, scheduler, save_model_path, n_epochs=0, ): """ Args: model: train_loader: valid_loader: criterion: optimiser: scheduler: save_model_path: n_epochs: Returns: """ valid_loss_min = numpy.Inf # track change in validation loss E = tqdm(range(1, n_epochs + 1)) for epoch in E: train_loss = 0.0 valid_loss = 0.0 dice_score = 0.0 with TorchTrainSession(model): train_set = tqdm(train_loader, postfix={"train_loss": 0.0}) for data, target in train_set: data, target = ( data.to(global_torch_device()), target.to(global_torch_device()), ) optimiser.zero_grad() output, *_ = model(data) output = torch.sigmoid(output) loss = criterion(output, target) loss.backward() optimiser.step() train_loss += loss.item() * data.size(0) train_set.set_postfix(ordered_dict={"train_loss": loss.item()}) with TorchEvalSession(model): with torch.no_grad(): validation_set = tqdm(valid_loader, postfix={ "valid_loss": 0.0, "dice_score": 0.0 }) for data, target in validation_set: data, target = ( data.to(global_torch_device()), target.to(global_torch_device()), ) # forward pass: compute predicted outputs by passing inputs to the model output, *_ = model(data) output = torch.sigmoid(output) # calculate the batch loss loss = criterion(output, target) # update average validation loss valid_loss += loss.item() * data.size(0) dice_cof = intersection_over_union( output.cpu().detach().numpy(), target.cpu().detach().numpy()) dice_score += dice_cof * data.size(0) validation_set.set_postfix(ordered_dict={ "valid_loss": loss.item(), "dice_score": dice_cof }) # calculate average losses train_loss = train_loss / len(train_loader.dataset) valid_loss = valid_loss / len(valid_loader.dataset) dice_score = dice_score / len(valid_loader.dataset) # print training/validation statistics E.set_description(f"Epoch: {epoch}" f" Training Loss: {train_loss:.6f} " f"Validation Loss: {valid_loss:.6f} " f"Dice Score: {dice_score:.6f}") # save model if validation loss has decreased if valid_loss <= valid_loss_min: print( f"Validation loss decreased ({valid_loss_min:.6f} --> {valid_loss:.6f}). Saving model ..." ) torch.save(model.state_dict(), save_model_path) valid_loss_min = valid_loss scheduler.step() model, scheduler = reschedule(model, epoch, scheduler) return model
def main(): dataset_root = Path.home() / "Data" base_path = ensure_existence(PROJECT_APP_PATH.user_data / 'maskrcnn') log_path = ensure_existence(PROJECT_APP_PATH.user_log / 'maskrcnn') export_root = ensure_existence(base_path / 'models') model_name = f'maskrcnn_pennfudanped' batch_size = 4 num_epochs = 10 optimiser_spec = GDKC(torch.optim.Adam, lr=3e-4) scheduler_spec = GDKC( torch.optim.lr_scheduler. StepLR, # a learning rate scheduler which decreases the learning rate by step_size=3, # 10x every 3 epochs gamma=0.1, ) num_workers = os.cpu_count() torch_seed(3825) dataset = PennFudanDataset(dataset_root / "PennFudanPed", Split.Training, return_variant=ReturnVariant.all) dataset_validation = PennFudanDataset( dataset_root / "PennFudanPed", Split.Validation, return_variant=ReturnVariant.all, ) split = SplitIndexer(len(dataset), validation=0.3, testing=0) split_indices = torch.randperm(split.total_num).tolist() data_loader = DataLoader( Subset(dataset, split_indices[:-split.validation_num]), batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_batch_fn, ) data_loader_val = DataLoader( Subset(dataset_validation, split_indices[-split.validation_num:]), batch_size=1, shuffle=False, num_workers=num_workers, collate_fn=collate_batch_fn, ) model = get_pretrained_instance_segmentation_maskrcnn( dataset.response_channels) optimiser = optimiser_spec(trainable_parameters(model)) lr_scheduler = scheduler_spec(optimiser) if True: model = load_model(model_name=model_name, model_directory=export_root) if True: with TorchTrainSession(model): with TensorBoardPytorchWriter(log_path / model_name) as writer: for epoch_i in tqdm(range(num_epochs), desc="Epoch #"): maskrcnn_train_single_epoch(model=model, optimiser=optimiser, data_loader=data_loader, writer=writer) lr_scheduler.step() # update the learning rate maskrcnn_evaluate( model, data_loader_val, writer=writer ) # evaluate on the validation dataset save_model(model, model_name=model_name, save_directory=export_root) if True: with TorchEvalSession(model): # put the model in evaluation mode img, _ = dataset_validation[ 0] # pick one image from the test set with torch.no_grad(): prediction = model([img.to(global_torch_device())]) from matplotlib import pyplot pyplot.imshow( Image.fromarray( img.mul(255).permute(1, 2, 0).byte().numpy())) pyplot.show() import cv2 pyplot.imshow( Image.fromarray(prediction[0]["masks"][0, 0].mul( 255).byte().cpu().numpy())) pyplot.show() (boxes, labels, scores) = ( prediction[0]["boxes"].to('cpu').numpy(), prediction[0]["labels"].to('cpu').numpy(), torch.sigmoid(prediction[0]["scores"]).to('cpu').numpy(), ) from draugr.opencv_utilities import draw_bounding_boxes from draugr.torch_utilities.images.conversion import quick_to_pil_image indices = scores > 0.1 cv2.namedWindow(model_name, cv2.WINDOW_NORMAL) cv2.imshow( model_name, draw_bounding_boxes( quick_to_pil_image(img), boxes[indices], labels=labels[indices], scores=scores[indices], #categories=categories, )) cv2.waitKey()
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
def run_webcam_demo( cfg: NOD, categories: Sequence[str], model_checkpoint: Path, score_threshold: float = 0.5, window_name: str = "SSD", ): """ :param categories: :type categories: :param cfg: :type cfg: :param model_checkpoint: :type model_checkpoint: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype:""" cpu_device = torch.device("cpu") transforms = SSDTransform(cfg.input.image_size, cfg.input.pixel_mean, split=SplitEnum.testing) model = SingleShotDetection(cfg) checkpointer = CheckPointer(model, save_dir=ensure_existence( PROJECT_APP_PATH.user_data / "results")) checkpointer.load(model_checkpoint, use_latest=model_checkpoint is None) print( f"Loaded weights from {model_checkpoint if model_checkpoint else checkpointer.get_checkpoint_file()}" ) model.post_init() model.to(global_torch_device()) with TorchEvalSession(model): for infos in tqdm(DictUnityEnvironment(connect_to_running=True)): info = next(iter(infos.values())) new_images = extract_all_cameras(info) image = next(iter(new_images.values()))[..., :3][..., ::-1] image = gamma_correct_float_to_byte(image) result = model( transforms(image)[0].unsqueeze(0).to(global_torch_device()))[0] height, width, *_ = image.shape result["boxes"][:, 0::2] *= width / result["img_width"] result["boxes"][:, 1::2] *= height / result["img_height"] (boxes, labels, scores) = ( result["boxes"].to(cpu_device).numpy(), result["labels"].to(cpu_device).numpy(), result["scores"].to(cpu_device).numpy(), ) indices = scores > score_threshold if show_image( draw_bounding_boxes( image, boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, ).astype(numpy.uint8), window_name, wait=1, ): break # esc to quit
def run_traced_webcam_demo( input_cfg: NOD, categories: List, score_threshold: float = 0.7, window_name: str = "SSD", onnx_exported: bool = False, ): """ :param onnx_exported: :type onnx_exported: :param input_cfg: :type input_cfg: :param categories: :type categories: :param score_threshold: :type score_threshold: :param window_name: :type window_name: :return: :rtype: """ pass import torch cpu_device = torch.device("cpu") transforms = SSDTransform(input_cfg.image_size, input_cfg.pixel_mean, split=Split.Testing) model = None if onnx_exported: import onnx onnx_model = onnx.load("torch_model.onnx") onnx.checker.check_model(onnx_model) import onnxruntime ort_session = onnxruntime.InferenceSession("torch_model.onnx") def to_numpy(tensor): return (tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()) x = None # compute onnxruntime output prediction ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(x)} ort_outs = ort_session.run(None, ort_inputs) else: import torch import io torch.jit.load("torch_model.traced") with open("torch_model.traced", "rb") as f: # Load ScriptModule from io.BytesIO object buffer = io.BytesIO(f.read()) model = torch.jit.load( buffer) # Load all tensors to the original device """ buffer.seek(0) torch.jit.load(buffer, map_location=torch.device('cpu')) # Load all tensors onto CPU, using a device buffer.seek(0) model = torch.jit.load(buffer, map_location='cpu') # Load all tensors onto CPU, using a string # Load with extra files. extra_files = torch._C.ExtraFilesMap() extra_files['foo.txt'] = 'bar' torch.jit.load('torch_model.traced', _extra_files=extra_files) print(extra_files['foo.txt']) #exit(0) """ with TorchDeviceSession( device=global_torch_device(cuda_if_available=False), model=model): with TorchEvalSession(model): for image in tqdm(frame_generator(cv2.VideoCapture(0))): result = SSDOut(*model( transforms(image)[0].unsqueeze(0).to( global_torch_device()))) height, width, *_ = image.shape result.boxes[:, 0::2] *= width / result.img_width.cpu().item() result.boxes[:, 1::2] *= height / result.img_height.cpu().item() (boxes, labels, scores) = ( result.boxes.to(cpu_device).numpy(), result.labels.to(cpu_device).numpy(), result.scores.to(cpu_device).numpy(), ) indices = scores > score_threshold cv2.namedWindow(window_name, cv2.WINDOW_NORMAL) cv2.imshow( window_name, draw_bounding_boxes( image, boxes[indices], labels=labels[indices], scores=scores[indices], categories=categories, score_font=ImageFont.truetype( PACKAGE_DATA_PATH / "Lato-Regular.ttf", 24, ), ).astype(numpy.uint8), ) if cv2.waitKey(1) == 27: break # esc to quit
def train_siamese( model: Module, optimiser: Optimizer, criterion: callable, *, writer: Writer = MockWriter(), train_number_epochs: int, data_dir: Path, train_batch_size: int, model_name: str, save_path: Path, save_best: bool = False, img_size: Tuple[int, int], validation_interval: int = 1, ): """ :param img_size: :type img_size: :param validation_interval: :type validation_interval: :param data_dir: :type data_dir: :param optimiser: :type optimiser: :param criterion: :type criterion: :param writer: :type writer: :param model_name: :type model_name: :param save_path: :type save_path: :param save_best: :type save_best: :param model: :type model: :param train_number_epochs: :type train_number_epochs: :param train_batch_size: :type train_batch_size: :return: :rtype: """ train_dataloader = DataLoader( PairDataset( data_path=data_dir, transform=transforms.Compose([ transforms.Grayscale(), transforms.Resize(img_size), transforms.ToTensor(), ]), split=Split.Training, ), shuffle=True, num_workers=4, batch_size=train_batch_size, ) valid_dataloader = DataLoader( PairDataset( data_path=data_dir, transform=transforms.Compose([ transforms.Grayscale(), transforms.Resize(img_size), transforms.ToTensor(), ]), split=Split.Validation, ), shuffle=True, num_workers=4, batch_size=train_batch_size, ) best = math.inf E = tqdm(range(0, train_number_epochs)) batch_counter = count() for epoch in E: for tss in train_dataloader: batch_i = next(batch_counter) with TorchTrainSession(model): o = [t.to(global_torch_device()) for t in tss] optimiser.zero_grad() loss_contrastive = criterion(model(*o[:2]), o[2].to(dtype=torch.float)) loss_contrastive.backward() optimiser.step() train_loss = loss_contrastive.cpu().item() writer.scalar("train_loss", train_loss, batch_i) if batch_counter.__next__() % validation_interval == 0: with TorchEvalSession(model): for tsv in valid_dataloader: ov = [t.to(global_torch_device()) for t in tsv] v_o, fact = model(*ov[:2]), ov[2].to(dtype=torch.float) valid_loss = criterion(v_o, fact).cpu().item() valid_accuracy = (accuracy(distances=v_o, is_diff=fact).cpu().item()) writer.scalar("valid_loss", valid_loss, batch_i) if valid_loss < best: best = valid_loss print(f"new best {best}") writer.blip("new_best", batch_i) if save_best: save_model_parameters( model, optimiser=optimiser, model_name=model_name, save_directory=save_path, ) E.set_description( f"Epoch number {epoch}, Current train loss {train_loss}, valid loss {valid_loss}, valid_accuracy {valid_accuracy}" ) return model