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)