def forward(self, noise_init, noise_amp, mode='rand'):
x_prev_out = self.body[0](F.pad(noise_init, self.p3d))
for idx, block in enumerate(self.body[1:], 1):
x_prev_out = torch.tanh(x_prev_out)
# Upscale
x_prev_out_up = utils.upscale(x_prev_out, idx, self.opt)
# Add noise if "random" sampling, else, add no noise is "reconstruction" mode
if mode == 'rand':
x_prev_out_up_2 = utils.interpolate_3D(
x_prev_out,
size=[
x_prev_out_up.shape[-3] + (self.opt.num_layer + 2) * 2,
x_prev_out_up.shape[-2] + (self.opt.num_layer + 2) * 2,
x_prev_out_up.shape[-1] + (self.opt.num_layer + 2) * 2
])
noise = utils.generate_noise(ref=x_prev_out_up_2)
x_prev = block(x_prev_out_up_2 + noise * noise_amp[idx])
else:
x_prev = block(F.pad(x_prev_out_up, self.p3d))
x_prev_out = x_prev + x_prev_out_up
out = torch.tanh(x_prev_out)
return out
def movnect(x, j):
with tf.compat.v1.variable_scope('movnect') as vs:
with tf.compat.v1.variable_scope('block13_a'):
init_x = x
x = slim.conv2d(x, 368, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.separable_conv2d(x, 368, 3, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.conv2d(x, 256, 1, 1, activation_fn=None, data_format='NCHW', normalizer_fn=slim.batch_norm)
init_x = slim.conv2d(init_x, 256, 1, 1, activation_fn=None, data_format='NCHW')
x = x + init_x
with tf.compat.v1.variable_scope('block13_b'):
x = slim.conv2d(x, 192, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.separable_conv2d(x, 192, 3, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.conv2d(x, 192, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
with tf.compat.v1.variable_scope('upsample_conv'):
x = utils.upscale(x, 2, data_format='NCHW')
non_j_x = slim.conv2d(x, 128, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
j_x = slim.conv2d(x, 3 * j, 1, 1, activation_fn=None, data_format='NCHW', normalizer_fn=slim.batch_norm)
dx, dy, dz = tf.split(j_x, 3, axis=1)
d_j = tf.abs(dx) + tf.abs(dy) + tf.abs(dz)
x = tf.concat([non_j_x, d_j, j_x], axis=1)
with tf.compat.v1.variable_scope('lsat'):
x = slim.conv2d(x, 128, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.separable_conv2d(x, 128, 3, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.separable_conv2d(x, 4*j, 1, 1, activation_fn=None, data_format='NCHW')
var = tf.contrib.framework.get_variables(vs)
return x, var # [256^2 * 4*j]
def vnect(x, j):
with tf.compat.v1.variable_scope('vnect') as vs:
with tf.compat.v1.variable_scope('block1_a'):
init_x = x
x = slim.conv2d(x, 512, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.conv2d(x, 512, 3, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.conv2d(x, 1024, 1, 1, activation_fn=None, data_format='NCHW', normalizer_fn=slim.batch_norm)
init_x = slim.conv2d(init_x, 1024, 1, 1, activation_fn=None, data_format='NCHW')
x = x + init_x
with tf.compat.v1.variable_scope('block1_b'):
x = slim.conv2d(x, 256, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.separable_conv2d(x, 128, 3, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.conv2d(x, 256, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
with tf.compat.v1.variable_scope('upsample_conv'):
x = utils.upscale(x, 2, 'NCHW')
non_j_x = slim.conv2d(x, 128, 4, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
j_x = slim.conv2d(x, 3 * j, 4, 1, activation_fn=None, data_format='NCHW', normalizer_fn=slim.batch_norm)
dx, dy, dz = tf.split(j_x, 3, axis=1)
d_j = tf.sqrt(tf.maximum(dx*dx+dy*dy+dz*dz, 0.))
# d_j = tf.concat([dx, dy, dz], axis=1) # temptemptemptemptemp
x = tf.concat([non_j_x, d_j, j_x], axis=1)
with tf.compat.v1.variable_scope('lsat'):
x = slim.conv2d(x, 128, 1, 1, activation_fn=tf.nn.relu, data_format='NCHW', normalizer_fn=slim.batch_norm)
x = slim.conv2d(x, 4*j, 1, 1, activation_fn=None, data_format='NCHW')
var = tf.contrib.framework.get_variables(vs)
return x, var # [256^2 * 4*j]
def dataset(dataset_train_path,batch_size,scale_factor):
assert(os.path.exists(dataset_train_path))
data = []
for file in os.listdir(dataset_train_path):
if file.endswith('.bmp'):
filepath = os.path.join(dataset_train_path,file)
img = imageio.imread(filepath).dot([0.299, 0.587, 0.114])
patches = extract_patches(img,(36,36),0.166)
data += [patches[idx] for idx in range(patches.shape[0])]
mod_data = [from_numpy(np.expand_dims(blur(upscale(patch,scale_factor),scale_factor),0)).float() for patch in data]
data = [from_numpy(np.expand_dims(upscale(patch,scale_factor),0)).float() for patch in data]
l = len(data)
for idx in range(0,l,batch_size):
yield stack(mod_data[idx:min(idx+batch_size,l)]),stack(data[idx:min(idx+batch_size,l)])
def save_solution(self, file_name="solution.png", scale=3):
"""Saves the solution as png."""
if self.solution is None:
raise util.MazeError(
"Cannot save solution because it is not assigned.\n"
"Use the \"solve\" method to solve a maze."
)
Image.fromarray(util.upscale(self.solution, scale), "RGB").save(file_name, "png")
def save_maze(self, file_name="maze.png", scale=3):
"""Saves the maze as png."""
if self.maze is None:
raise util.MazeError(
"Cannot save maze because it is not assigned.\n"
"Use the \"create\" or \"load_maze\" method to create or load a maze."
)
Image.fromarray(util.upscale(self.maze, scale), "RGB").save(file_name, "png")
def save_maze(self, file_name="maze.png", line=5, path=15):
"""Saves the maze as png."""
if self.maze is None:
raise util.MazeError(
"Cannot save maze because it is not assigned.\n"
"Use the \"create\" or \"load_maze\" method to create or load a maze."
)
img = Image.fromarray(util.upscale(self.maze, line, path), "RGB")
img.save("uncropped.png")
print(img)
img = PIL.ImageOps.invert(img).convert("L")
newImage = Image.new("L", (1085, 720))
newImage.paste(img, (0, 7, 1085, 7 + 705))
newImage.save(file_name, "png")
return newImage
def refinement_layers(self, start_idx, x_prev_out, noise_amp, mode):
for idx, block in enumerate(self.body[start_idx:], start_idx):
if self.opt.vae_levels == idx + 1:
x_prev_out.detach_()
# Upscale
x_prev_out_up = utils.upscale(x_prev_out, idx + 1, self.opt)
# Add noise if "random" sampling, else, add no noise is "reconstruction" mode
if mode == 'rand':
noise = utils.generate_noise(ref=x_prev_out_up)
x_prev = block(x_prev_out_up + noise * noise_amp[idx + 1])
else:
x_prev = block(x_prev_out_up)
x_prev_out = torch.tanh(x_prev + x_prev_out_up)
return x_prev_out
def train_and_validate(dataset_path,
batch_size,
scale_factor,
num_epochs,
learning_rate,
weight_decay,
output_dir,
verbose=True):
model_output_dir = os.path.join(output_dir, 'model')
model = network.ten()
logging.info('TEN Model loaded')
if verbose:
print('TEN Model loaded')
total_params = sum(p.numel() for p in model.parameters())
print(f'Total Parameters: {total_params}')
if cuda:
model = model.cuda()
model.train()
optimizer = torch.optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=weight_decay)
logging.info('Adam optimizer loaded')
if verbose:
print('Adam optimizer loaded')
total_epoch_time = 0
losses = []
# Set initial best_loss arbitrarily high
best_val_loss = 2.0e50
for epoch in range(1, num_epochs + 1):
model.train()
logging.info(f'Epoch {epoch} Start')
if verbose:
print(f'\n<----- START EPOCH {epoch} ------->\n')
start_time = time()
total_loss_for_this_epoch = 0
# For each batch
batch = 1
for patches, gt in dataset(dataset_path, batch_size, scale_factor):
optimizer.zero_grad()
if cuda:
patches = patches.cuda()
gt = gt.cuda()
pred = model(patches)
loss = compute_loss(pred, gt)
logging.info(f'Epoch {epoch} Batch {batch} Loss {loss.item()}')
if verbose and (batch - 1) % 10 == 0:
print(f'Epoch {epoch} Batch {batch} Loss {loss.item()}')
loss.backward()
optimizer.step()
total_loss_for_this_epoch += loss.item()
batch += 1
avg_loss = total_loss_for_this_epoch / batch
losses.append(avg_loss)
epoch_time = time() - start_time
if verbose:
print(f'Epoch time: {epoch_time}')
total_epoch_time += epoch_time
# Validation
model.eval()
val_img_file = random.choice(
[f for f in os.listdir(dataset_path) if f.endswith('.bmp')])
val_img = imageio.imread(os.path.join(dataset_path, val_img_file)).dot(
[0.299, 0.587, 0.114])
mod_val_img = torch.from_numpy(
blur(upscale(val_img, scale_factor),
scale_factor)).float().unsqueeze(0).unsqueeze(0)
val_img = torch.from_numpy(upscale(
val_img, scale_factor)).float().unsqueeze(0).unsqueeze(0)
if cuda:
mod_val_img = mod_val_img.cuda()
val_img = val_img.cuda()
out = model(mod_val_img)
val_loss = compute_loss(out, val_img).item()
if verbose:
print(
f'Epoch {epoch} Validation Image {val_img_file} Loss {val_loss}'
)
# Save current model
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, model_output_dir, 'current.pth')
logging.info('Current model saved')
if verbose:
print('Current model saved')
# Save best model
if val_loss < best_val_loss:
best_val_loss = val_loss
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, model_output_dir, 'best.pth')
logging.info('Best model saved')
if verbose:
print('Best model saved')
# Save model every 20 epochs
if (epoch) % 20 == 0:
save_checkpoint(
{
'epoch': epoch,
'state_dict': model.state_dict(),
'optimizer': optimizer.state_dict(),
}, model_output_dir, f'epoch_{epoch}.pth')
logging.info(f'Epoch {epoch} Model saved')
if verbose:
print(f'Epoch {epoch} Model saved')
# Learning rate decay
if epoch % 30 == 0 and epoch <= 60:
learning_rate = learning_rate / 10
for param_group in optimizer.param_groups:
param_group['lr'] = learning_rate
logging.info(
f'Epoch {epoch}: Learning rate decayed by factor of 10')
logging.info(f'Epoch {epoch} completed')
if verbose:
print(
f'\n<----- END EPOCH {epoch} Time elapsed: {time()-start_time}------->\n'
)
logging.info('All epochs completed')
logging.info(f'Average Time: {total_epoch_time/num_epochs:.4f} seconds')
logging.info(f'Average Loss: {sum(losses) / len(losses):.4f}')
if verbose:
print('All epochs completed')
print(f'Average Time: {total_epoch_time/num_epochs:.4f} seconds')
print(f'Average Loss: {sum(losses) / len(losses):.4f}')
if verbose:
print('Losses array: ', losses)
print('Best Validation Loss', best_val_loss)