def train(self, data_loader, epoch, history, **kwargs): progress = tqdm.tqdm(data_loader) for i, (source, target) in enumerate(progress): source, target = to_var(source), to_var(target) output = self.gnet(source) d_loss, d_terms = self.foward_d(source, output.detach(), target) g_loss, g_terms = self.foward_g(source, output, target) acc = self.metric(output, target) self.d_optim.zero_grad() d_loss.backward() self.d_optim.step() self.g_optim.zero_grad() g_loss.backward() self.g_optim.step() progress.set_description('Epoch#%d' % (epoch + 1)) progress.set_postfix(history.add({**g_terms, **d_terms, 'acc/psnr': acc})) self.logger.image( {'input': source.data, 'output': output.data, 'target': target.data}, epoch=epoch, prefix='train_') self.logger.scalar(history.metric(), epoch)
def test_semantic_segmentation_confusion(): num_class = 7 gt = to_var(torch.LongTensor(4, 3).random_(to=num_class)) pred = to_var(torch.LongTensor(4, 3).random_(to=num_class)) confusion = onegan.metrics.semantic_segmentation.confusion_table(pred, gt, num_class) assert confusion.sum() == 4 * 3
def test_semantic_segmentation_iou(): num_class = 7 gt = to_var(torch.LongTensor(4, 3).random_(0, to=2)) pred = to_var(torch.LongTensor(4, 3).random_(3, to=4)) confusion = onegan.metrics.semantic_segmentation.confusion_table(pred, gt, num_class) iou = onegan.metrics.semantic_segmentation.intersection_over_union(confusion) assert len(iou) == num_class
def test_semantic_segmentation_metric(): num_class = 3 gt = to_var(torch.LongTensor(3, 3).random_(to=num_class)) pred = to_var(torch.LongTensor(3, 3).random_(to=num_class)) metric = onegan.metrics.semantic_segmentation.Metric(num_class=num_class, only_scalar=True, prefix='acc/') scalar_result = metric(pred, gt) metric = onegan.metrics.semantic_segmentation.Metric(num_class=num_class) full_result = metric(pred, gt) assert scalar_result.keys() != full_result.keys()
def test_to_var(): dummy_tensor = torch.rand(8, 7) if torch.cuda.is_available(): assert hasattr(to_var(dummy_tensor), 'get_device') set_device_mode('gpu') assert hasattr(to_var(dummy_tensor), 'get_device') set_device_mode('cpu') with pytest.raises(Exception): to_var(dummy_tensor).get_device()
def test_loss_terms(): source = to_device(torch.FloatTensor(10, 3, 128, 128)) target = to_device(torch.FloatTensor(10, 3, 128, 128)) source, target = to_var(source), to_var(target) output = g(source) real = losses.conditional_input(source, target, conditional) fake = losses.conditional_input(source, output, conditional) losses.adversarial_ce_loss(F.sigmoid(d(fake)), 1) losses.adversarial_ls_loss(d(fake), 1) losses.adversarial_w_loss(d(fake), True) losses.gradient_penalty(d, real, fake)
def _closure(model, data, volatile=False): source = to_var(data['image'], volatile=volatile) target = to_var(data['label'], volatile=volatile) score, pred_type = model(source) _, output = torch.max(score, 1) set_gallery.image(merge_viz(data['image'], data['label'], output.data.cpu()), filenames=data['filename']) gallery.image(output.data.float() * (255 / 5), filenames=data['filename']) accuracy = metric(output, target) accuracy['score'] = score_metric(output, target) return accuracy
def objective(score, prediction, pred_type, target, data): ''' per-pixel classification loss ''' if args.focal_gamma: seg_loss = focal_loss(score, target) else: seg_loss = ce_loss(score, target) terms = {'loss/loss': seg_loss, 'loss/cla': seg_loss} ''' area smoothness loss ''' if args.l1_factor or args.l2_factor: l_loss = F.mse_loss if args.l2_factor else F.l1_loss l1_λ = args.l1_factor or args.l2_factor onehot_target = torch.zeros_like(score).scatter_( 1, target.unsqueeze(1), 1) l1_loss = l_loss(score, onehot_target) terms['loss/loss'] += l1_loss * l1_λ terms['loss/area'] = l1_loss ''' layout edge constraint loss ''' if args.edge_factor: edge_map = layout_gradient(prediction) edge_loss = F.binary_cross_entropy(edge_map, to_var(data['edge'])) terms['loss/loss'] += edge_loss * args.edge_factor terms['loss/edge'] = edge_loss ''' room type ce loss ''' # if args.type_λ: # type_loss = F.cross_entropy(pred_type.squeeze(), to_var(data['type'].long())) # terms['loss/loss'] += type_loss * args.type_λ # terms['loss/room_type'] = type_loss return terms
def evaluate(self, data_loader, epoch, history, **kwargs): progress = tqdm.tqdm(data_loader, leave=False) for i, (source, target) in enumerate(progress): source, target = to_var(source, volatile=True), to_var(target, volatile=True) output = self.gnet(source) _, d_terms = self.foward_d(source, output.detach(), target) _, g_terms = self.foward_g(source, output, target) acc = self.metric(output, target) progress.set_description('Evaluate') progress.set_postfix(history.add({**g_terms, **d_terms, 'acc/psnr': acc}, log_suffix='_val')) self.logger.image( {'input': source.data, 'output': output.data, 'target': target.data}, epoch=epoch, prefix='val_') self.logger.scalar(history.metric(), epoch)
def _closure(g, d, data, volatile=False): AtoB = args.which_direction == 'AtoB' source = to_var(data['A' if AtoB else 'B'], volatile=volatile) target = to_var(data['B' if AtoB else 'A'], volatile=volatile) output = g(source) # fake fake = pair_pool.query(torch.cat((source, output), dim=1).data) pred_fake = d(fake.detach()) loss_d_fake = gan_loss(pred_fake, 0) # real real = torch.cat((source, target), dim=1) pred_real = d(real) loss_d_real = gan_loss(pred_real, 1) yield { 'loss/loss_d': (loss_d_fake + loss_d_real) * 0.5, 'loss/d_fake': loss_d_fake, 'loss/d_real': loss_d_real, } # generated fake = torch.cat([source, output], dim=1) pred_fake = d(fake) loss_g_gan = gan_loss(pred_fake, 1) loss_g_l1 = l1_loss(output, target) yield { 'loss/loss_g': loss_g_gan * args.lambda_A + loss_g_l1, 'loss/g_l1': loss_g_l1, 'loss/g_gan': loss_g_gan, } accuracy = {'acc/psnr': psnr(output, target)} yield accuracy viz_results = { 'input': source.data, 'output': output.data, 'target': target.data } tensorboard.image(viz_results, epoch=estimator.state['epoch'], prefix='val_' if volatile else 'train_') yield
def _closure(model, data, volatile=False): source = to_var(data['image'], volatile=volatile) target = to_var(data['label'], volatile=volatile) score, pred_type = model(source) _, output = torch.max(score, 1) loss = objective(score, output, pred_type, target, data) accuracy = metric(output, target) if volatile: accuracy['score'] = score_metric(output, target) viz_results = { 'input': source.data, 'output': output.data.float(), 'target': target.data.float() } tensorboard.image(viz_results, epoch=estimator.state['epoch'], prefix='val_' if volatile else 'train_') return loss, accuracy
def __init__(self, kernel, padding=1, dilation=1, name='VisionConv2d'): """ Args: kernel: str or np.ndarray """ if isinstance(kernel, str): kernel = { 'laplacian': laplacian_kernel, 'sobel_vertical': sobel_vertical_kernel, 'sobel_horizontal': sobel_horizontal_kernel, }[kernel]() assert kernel.ndim == 2, 'Plain Vision Kernel should be 2D' self.kernel = to_var(self._to_tensor(kernel[np.newaxis, np.newaxis, :])) self.padding = padding self.dilation = dilation self.name = name
def test_img_normalize(): dummy_img_tensor = torch.rand(3, 64, 64) transformed_tensor = dummy_img_tensor.add_(-0.5).div_(0.5) normalized_tensor = img_normalize(transformed_tensor) assert normalized_tensor.min() == 0 assert normalized_tensor.max() == 1 normalized_var = img_normalize(to_var(dummy_img_tensor)) assert normalized_var.data.min() == 0 assert normalized_var.data.max() == 1 dummy_tensor = torch.LongTensor(3, 8, 8).random_(-2, to=2).float() normalized_tensor = img_normalize(dummy_tensor, img_range=(-2, 2)) assert normalized_tensor.min() >= 0 assert normalized_tensor.max() <= 1 dummy_zero_tensor = torch.zeros(8, 8) img_normalize(dummy_zero_tensor)
def _closure(models, data, volatile=False): ga, gb, da, db = models AtoB = args.which_direction == 'AtoB' real_a = to_var(data['A' if AtoB else 'B'], volatile=volatile) real_b = to_var(data['B' if AtoB else 'A'], volatile=volatile) if args.identity > 0: idt_a = ga(real_b) loss_idt_a = F.l1_loss(idt_a, real_b) * args.lambda_B * args.identity idt_b = gb(real_a) loss_idt_b = F.l1_loss(idt_b, real_a) * args.lambda_A * args.identity else: loss_idt_a = loss_idt_b = 0 fake_a = gb(real_b) fake_b = ga(real_a) loss_g_a = gan_loss(da(fake_b), 1) loss_g_b = gan_loss(db(fake_a), 1) # forward cycle loss rec_a = gb(fake_b) loss_cycle_a = F.l1_loss(rec_a, real_a) * args.lambda_A # backward cycle loss rec_b = ga(fake_a) loss_cycle_b = F.l1_loss(rec_b, real_b) * args.lambda_B yield { 'loss/g': loss_g_a + loss_g_b + loss_cycle_a + loss_cycle_b + loss_idt_a + loss_idt_b, 'loss/g_a': loss_g_a, 'loss/g_b': loss_g_b, 'loss/cycle_a': loss_cycle_a, 'loss/cycle_b': loss_cycle_b, }, (optim_g, 'loss/g') def _d(d, real, fake): loss_d_real = gan_loss(d(real), 1) loss_d_fake = gan_loss(d(fake.detach()), 0) return (loss_d_real + loss_d_fake) * 0.5 loss_d_a = _d(da, real_b, b_pool.query(fake_b.data)) yield {'loss/d_a': loss_d_a}, (optim_da, 'loss/d_a') loss_d_b = _d(db, real_a, a_pool.query(fake_a.data)) yield {'loss/d_b': loss_d_b}, (optim_db, 'loss/d_b') viz_results = { 'realA': real_a.data, 'realB': real_b.data, 'fakeA': fake_a.data, 'fakeB': fake_b.data, 'recA': rec_a.data, 'recB': rec_b.data } if args.identity > 0: viz_results['idtA'] = idt_a viz_results['idtB'] = idt_b tensorboard.image(viz_results, epoch=estimator.state['epoch'], prefix='val_' if volatile else 'train_') yield
def test_psnr(): dummy_output = to_var(img_normalize(torch.rand(4, 3, 64, 64))) dummy_target = to_var(img_normalize(torch.rand(4, 3, 64, 64))) psnr = onegan.metrics.psnr(dummy_output, dummy_target) assert psnr
def _closure(model, data, volatile=False): source = to_var(data['image'], volatile=volatile) score, pred_type = model(source) _, output = torch.max(score, 1) score = score_metric(output, data['label']) return {'score': score}
def fetch_data(): source, target = next(progress) return to_var(source), to_var(target)