def test_train_steps(self):
# Get real and fake images
real_batch = common.load_images(self.N, size=self.H)
# Setup optimizers
optD = optim.Adam(self.netD.parameters(), 2e-4, betas=(0.0, 0.9))
optG = optim.Adam(self.netG.parameters(), 2e-4, betas=(0.0, 0.9))
# Log statistics to check
log_data = metric_log.MetricLog()
# Test D train step
log_data = self.netD.train_step(real_batch=real_batch,
netG=self.netG,
optD=optD,
device='cpu',
log_data=log_data)
log_data = self.netG.train_step(real_batch=real_batch,
netD=self.netD,
optG=optG,
log_data=log_data,
device='cpu')
for name, metric_dict in log_data.items():
assert type(name) == str
assert type(metric_dict['value']) == float
def test_load_and_save_image(self):
image, label = common.load_images(n=1)
image = torch.squeeze(image, dim=0)
output_file = os.path.join(self.log_dir, 'test_img.png')
common.save_tensor_image(image, output_file=output_file)
check = np.array(Image.open(output_file))
assert check.shape == (32, 32, 3)
assert label.shape == (1, )
def test_rot_tensor(self):
# Load image and model
image, _ = common.load_images(1, size=32)
# For any rotation, after performing the same action 4 times,
# you should return to the same pixel value
for deg in [0, 90, 180, 270]:
x = image.clone()
for _ in range(4):
x = self.netD._rot_tensor(x, deg)
assert torch.sum((x - image)**2) < 1e-5
def test_rotate_batch(self):
# Load image and model
images, _ = common.load_images(8, size=32)
check = images.clone()
check, labels = self.netD._rotate_batch(check)
degrees = [0, 90, 180, 270]
# Rotate 3 more times to get back to original.
for i in range(check.shape[0]):
for _ in range(3):
check[i] = self.netD._rot_tensor(check[i], degrees[labels[i]])
assert torch.sum((images - check)**2) < 1e-5
