def train(model, dset_loaders, optimizer, lr_scheduler, criterion, epoch, cuda, clip, writer): """ Trains model using data_loader with the given optimizer, lr_scheduler, criterion and epoch """ lr_scheduler.step() # Each epoch has a training and validation phase for phase in ['train', 'val']: since = time.time() if phase == 'train': model.train() # Set model to training mode else: model.eval() # Set model to evaluate mode running_loss = 0.0 for iteration, batch in enumerate(dset_loaders[phase]): # Use this if doing cyclic learning # lr_scheduler.batch_step() targets = batch['targets'] inputs = batch['inputs'] inputs.requires_grad_() targets.requires_grad_(False) if cuda: inputs = inputs.cuda() targets = targets.cuda() # forward if iteration == 0: print("Loaded " + phase + " batch in {:.0f}s".format(time.time() - since)) residuals = model(inputs) outputs = inputs + residuals outputs = torch.clamp(outputs, 0.0, 1.0) loss = criterion(outputs, targets) optimizer.zero_grad() # backward + optimize only if in training phase if phase == 'train': loss.backward() nn.utils.clip_grad_norm_(model.parameters(), clip) optimizer.step() # statistics running_loss += loss.item() if iteration == 0 and cuda: cnn_utils.print_mem_usage() if iteration % 100 == 0: print("===> Epoch[{}]({}/{}): Loss: {:.5f}".format( epoch, iteration, len(dset_loaders[phase]), loss.item())) if phase == 'train': if not cnn_utils.check_gradients(model): print( "No gradients are being computed during training") exit(-1) if iteration == len(dset_loaders[phase]) - 1: desired_shape = [int(shape[0]) for shape in batch['shape']] inputs_s = undo_remap(inputs[0], desired_shape, dtype=torch.float32) residuals_s = undo_remap(residuals[0], desired_shape, dtype=torch.float32) outputs_s = undo_remap(outputs[0], desired_shape, dtype=torch.float32) targets_s = undo_remap(targets[0], desired_shape, dtype=torch.float32) input_imgs = cnn_utils.transform_lf_to_torch(inputs_s) residual_imgs = cnn_utils.transform_lf_to_torch(residuals_s) out_imgs = cnn_utils.transform_lf_to_torch(outputs_s) truth_imgs = cnn_utils.transform_lf_to_torch(targets_s) input_grid = vutils.make_grid(input_imgs, nrow=8, range=(0, 1), normalize=True, pad_value=1.0) residual_grid = vutils.make_grid(residual_imgs, nrow=8, range=(-1, 1), normalize=True, pad_value=1.0) output_grid = vutils.make_grid(out_imgs, nrow=8, range=(0, 1), normalize=True, pad_value=1.0) target_grid = vutils.make_grid(truth_imgs, nrow=8, range=(0, 1), normalize=True, pad_value=1.0) diff_grid = vutils.make_grid(torch.abs(truth_imgs - out_imgs), nrow=8, range=(0, 1), normalize=True, pad_value=1.0) writer.add_image(phase + '/input', input_grid, epoch) writer.add_image(phase + '/residual', residual_grid, epoch) writer.add_image(phase + '/output', output_grid, epoch) writer.add_image(phase + '/target', target_grid, epoch) writer.add_image(phase + '/difference', diff_grid, epoch) epoch_loss = running_loss / len(dset_loaders[phase]) writer.add_scalar(phase + '/loss', epoch_loss, epoch) print("Phase {} average overall loss {:.5f}".format(phase, epoch_loss)) time_elapsed = time.time() - since print("Phase {} took {:.0f}s overall".format(phase, time_elapsed)) if phase == 'val': print() for idx, param_group in enumerate(optimizer.param_groups): writer.add_scalar('learning_rate', param_group['lr'], epoch) return epoch_loss
def process(self): """Perform the model warping and output an image grid""" if self.getPropertyByIdentifier("off").value: print("Image warping is currently turned off") return 1 start_time = time.time() if self.getPropertyByIdentifier("display_input").value: im_data = [] for name in INPORT_LIST: im_data.append(self.getInport(name).getData()) out_image = Image(OUT_SIZE, DTYPE) out = resize( im_data[0].colorLayers[0].data.transpose(1, 0, 2), OUT_SIZE_LIST) with warnings.catch_warnings(): warnings.simplefilter("ignore") inter_out = img_as_ubyte(out) out_image.colorLayers[0].data = inter_out self.getOutport("outport").setData(out_image) return 1 if model is None: print("No model for synthesis") return -1 cam = inviwopy.app.network.EntryExitPoints.camera im_data = [] for name in INPORT_LIST: im_data.append(self.getInport(name).getData()) for im in im_data: if not (im_data[0].dimensions == im.dimensions): print("Operation is incompatible with images of different size") print("Size 1: ", im_data[0].dimensions) print("Size 2: ", im.dimensions) return -1 out_image = Image(OUT_SIZE, DTYPE) sample_image = Image(SAMPLE_SIZE, DTYPE) im_colour = [] for idx, name in enumerate(INPORT_LIST): im_colour.append(im_data[idx].colorLayers[0].data[:, :, :3].transpose(1, 0, 2)) im_depth = [] near = cam.nearPlane far = cam.farPlane baseline = 0.5 focal_length = cam.projectionMatrix[0][0] fov = cam.fov.value for idx, name in enumerate(INPORT_LIST): im_depth.append( conversions.depth_to_pixel_disp( im_data[idx].depth.data.transpose(1, 0), near=near, far=far, baseline=baseline, focal_length=focal_length, fov=fov, image_pixel_size=float(im_data[0].dimensions[0])) ) sample = { 'depth': torch.tensor(im_depth[0], dtype=torch.float32).unsqueeze_(0), 'colour': torch.tensor(im_colour[0], dtype=torch.float32).unsqueeze_(0), 'grid_size': GRID_SIZE} warped = data_transform.transform_inviwo_to_warped(sample) desired_shape = warped['shape'] im_input = warped['inputs'].unsqueeze_(0) if cuda: im_input = im_input.cuda() model.eval() output = model(im_input) output += im_input output = torch.clamp(output, 0.0, 1.0) end_time = time.time() - start_time print("Grid light field rendered in {:4f}".format(end_time)) out_unstack = data_transform.undo_remap( output[0], desired_shape, dtype=torch.float32) out_colour = cnn_utils.transform_lf_to_torch( out_unstack ) output_grid = vutils.make_grid( out_colour, nrow=8, range=(0, 1), normalize=False, padding=2, pad_value=1.0) output_grid = resize( output_grid.cpu().detach().numpy().transpose(1, 2, 0), OUT_SIZE_LIST) with warnings.catch_warnings(): warnings.simplefilter("ignore") inter_out = img_as_ubyte(output_grid) #inter_out = denormalise_lf(output_grid) #inter_out = inter_out.cpu().detach().numpy().astype(np.uint8).transpose(1, 2, 0) # Add an alpha channel here shape = tuple(OUT_SIZE_LIST) + (4,) final_out = np.full(shape, 255, np.uint8) final_out[:, :, :3] = inter_out shape = tuple(SAMPLE_SIZE_LIST) + (4,) sample_out = np.full(shape, 255, np.uint8) sample_out[:, :, :3] = np.around( data_transform.denormalise_lf( out_unstack).cpu().detach().numpy() ).astype(np.uint8)[self.getPropertyByIdentifier("sample_num").value] # Inviwo expects a uint8 here out_image.colorLayers[0].data = final_out sample_image.colorLayers[0].data = sample_out self.getOutport("outport").setData(out_image) self.getOutport("sample").setData(sample_image) end_time = time.time() - start_time print("Overall render time was {:4f}".format(end_time))
def train(model, dset_loaders, optimizer, lr_scheduler, criterion, epoch, cuda, clip, writer, schedule_type): """ Trains model using data_loader with the given optimizer, lr_scheduler, criterion and epoch """ if schedule_type == 'warm': lr_scheduler.step() # Each epoch has a training and validation phase for phase in ['train', 'val']: since = time.time() if phase == 'train': model.train() # Set model to training mode else: model.eval() # Set model to evaluate mode running_loss = 0.0 for iteration, batch in enumerate(dset_loaders[phase]): if schedule_type == 'cyclical': lr_scheduler.batch_step() targets = batch['targets'] inputs = batch['inputs'] inputs.requires_grad_() targets.requires_grad_(False) if cuda: inputs = inputs.cuda() targets = targets.cuda() # forward if iteration == 0: print("Loaded " + phase + " batch in {:.0f}s".format(time.time() - since)) outputs = model(inputs) loss = criterion(outputs, targets) optimizer.zero_grad() # backward + optimize only if in training phase if phase == 'train': loss.backward() nn.utils.clip_grad_norm_(model.parameters(), clip) optimizer.step() # statistics running_loss += loss.item() if iteration == 0 and cuda: cnn_utils.print_mem_usage() if iteration % 100 == 0: print("===> Epoch[{}]({}/{}): Loss: {:.5f}".format( epoch, iteration, len(dset_loaders[phase]), loss.item())) if iteration == len(dset_loaders[phase]) - 1: input_imgs = cnn_utils.transform_lf_to_torch(inputs[0]) out_imgs = cnn_utils.transform_lf_to_torch(outputs[0]) truth_imgs = cnn_utils.transform_lf_to_torch(targets[0]) input_grid = vutils.make_grid(input_imgs, nrow=8, range=(-1, 1), normalize=True, pad_value=1.0) output_grid = vutils.make_grid(out_imgs, nrow=8, range=(-1, 1), normalize=True, pad_value=1.0) target_grid = vutils.make_grid(truth_imgs, nrow=8, range=(-1, 1), normalize=True, pad_value=1.0) diff_grid = vutils.make_grid(torch.abs(truth_imgs - out_imgs), nrow=8, range=(-2, 2), normalize=True, pad_value=1.0) writer.add_image(phase + '/input', input_grid, epoch) writer.add_image(phase + '/output', output_grid, epoch) writer.add_image(phase + '/target', target_grid, epoch) writer.add_image(phase + '/difference', diff_grid, epoch) epoch_loss = running_loss / len(dset_loaders[phase]) writer.add_scalar(phase + '/loss', epoch_loss, epoch) print("Phase {} average overall loss {:.5f}".format(phase, epoch_loss)) time_elapsed = time.time() - since print("Phase {} took {:.0f}s overall".format(phase, time_elapsed)) if phase == 'val': print() if schedule_type == 'step': lr_scheduler.step() for idx, param_group in enumerate(optimizer.param_groups): writer.add_scalar('learning_rate', param_group['lr'], epoch) return epoch_loss