def reconstruct(random_sample):
n = 128
data, _ = get_batch(test_data, test_index[:n])
data = data.to(device)
if not colorful:
data = torch.ge(data, 0.5).float()
recon_img, _ = mae.reconstruct(data, random_sample=random_sample)
comparison = torch.cat([data, recon_img], dim=0).cpu()
reorder_index = torch.from_numpy(np.array([[i + j * n for j in range(2)] for i in range(n)])).view(-1)
comparison = comparison[reorder_index]
image_file = 'reconstruct.random.png' if random_sample else 'reconstruct.fixed.png'
if colorful:
save_image(comparison, os.path.join(result_path, image_file), nrow=16, normalize=True, scale_each=True, range=(-1, 1))
else:
save_image(comparison, os.path.join(result_path, image_file), nrow=16)
def reconstruct():
mae_shadow.eval()
n = 128
data, _ = get_batch(test_data, test_index[:n])
data = data.to(device)
recon_img, _ = mae_shadow.reconstruct(data)
comparison = torch.cat([data, recon_img], dim=0).cpu()
reorder_index = torch.from_numpy(
np.array([[i + j * n for j in range(2)] for i in range(n)])).view(-1)
comparison = comparison[reorder_index]
image_file = 'reconstruct.png'
save_image(comparison,
os.path.join(result_path, image_file),
nrow=16,
normalize=True,
scale_each=True,
range=(-1, 1))
def reconstruct():
mae_shadow.eval()
n = 128
data, _ = get_batch(test_data, test_index[:n])
data = data.to(device)
binary_data = torch.ge(data, 0.5).float()
recon_img, recon_probs = mae_shadow.reconstruct(binary_data)
comparison = torch.cat([data, recon_probs, binary_data, recon_img], dim=0).cpu()
reorder_index = torch.from_numpy(np.array([[i + j * n for j in range(4)] for i in range(n)])).view(-1)
comparison = comparison[reorder_index]
image_file = 'reconstruct.fixed.png'
save_image(comparison, os.path.join(result_path, image_file), nrow=32)
recon_img, recon_probs = mae_shadow.reconstruct(binary_data, random_sample=True)
comparison = torch.cat([data, recon_probs, binary_data, recon_img], dim=0).cpu()
reorder_index = torch.from_numpy(np.array([[i + j * n for j in range(4)] for i in range(n)])).view(-1)
comparison = comparison[reorder_index]
image_file = 'reconstruct.random.png'
save_image(comparison, os.path.join(result_path, image_file), nrow=32)
test_binary_index = np.arange(len(test_binary_data))
np.random.shuffle(val_binary_index)
np.random.shuffle(test_binary_index)
print(len(train_data))
print(len(val_binary_data))
print(len(test_binary_data))
polyak_decay = args.polyak
params = json.load(open(args.config, 'r'))
json.dump(params, open(os.path.join(model_path, 'config.json'), 'w'), indent=2)
mae = MAE.from_params(params).to(device)
# initialize
init_batch_size = 500
init_index = np.random.choice(train_index, init_batch_size, replace=False)
init_data, _ = get_batch(train_data, init_index)
init_data = binarize_image(init_data).to(device)
mae.eval()
mae.initialize(init_data, init_scale=1.0)
# create shadow mae for ema
params = json.load(open(args.config, 'r'))
mae_shadow = MAE.from_params(params).to(device)
exponentialMovingAverage(mae, mae_shadow, polyak_decay, init=True)
print(args)
def get_optimizer(learning_rate, parameters):
if opt == 'adam':
return optim.Adam(parameters, lr=learning_rate)
elif opt == 'adamax':
return optim.Adamax(parameters, lr=learning_rate)