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