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
