def save_img_results(imgs_tcpu, fake_imgs, num_imgs,
count, image_dir, summary_writer):
num = cfg.TRAIN.VIS_COUNT
# The range of real_img (i.e., self.imgs_tcpu[i][0:num])
# is changed to [0, 1] by function vutils.save_image
real_img = imgs_tcpu[-1][0:num]
vutils.save_image(
real_img, '%s/real_samples.png' % (image_dir),
normalize=True)
real_img_set = vutils.make_grid(real_img).numpy()
real_img_set = np.transpose(real_img_set, (1, 2, 0))
real_img_set = real_img_set * 255
real_img_set = real_img_set.astype(np.uint8)
sup_real_img = summary.image('real_img', real_img_set)
summary_writer.add_summary(sup_real_img, count)
for i in range(num_imgs):
fake_img = fake_imgs[i][0:num]
# The range of fake_img.data (i.e., self.fake_imgs[i][0:num])
# is still [-1. 1]...
vutils.save_image(
fake_img.data, '%s/count_%09d_fake_samples%d.png' %
(image_dir, count, i), normalize=True)
fake_img_set = vutils.make_grid(fake_img.data).cpu().numpy()
fake_img_set = np.transpose(fake_img_set, (1, 2, 0))
fake_img_set = (fake_img_set + 1) * 255 / 2
fake_img_set = fake_img_set.astype(np.uint8)
sup_fake_img = summary.image('fake_img%d' % i, fake_img_set)
summary_writer.add_summary(sup_fake_img, count)
summary_writer.flush()
def on_end_epoch(state):
print('[Epoch %d] Training Loss: %.4f (Accuracy: %.2f%%)' % (
state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0]))
train_loss_logger.log(state['epoch'], meter_loss.value()[0])
train_error_logger.log(state['epoch'], meter_accuracy.value()[0])
reset_meters()
engine.test(processor, get_iterator(False))
test_loss_logger.log(state['epoch'], meter_loss.value()[0])
test_accuracy_logger.log(state['epoch'], meter_accuracy.value()[0])
confusion_logger.log(confusion_meter.value())
print('[Epoch %d] Testing Loss: %.4f (Accuracy: %.2f%%)' % (
state['epoch'], meter_loss.value()[0], meter_accuracy.value()[0]))
torch.save(model.state_dict(), 'epochs/epoch_%d.pt' % state['epoch'])
# Reconstruction visualization.
test_sample = next(iter(get_iterator(False)))
ground_truth = (test_sample[0].unsqueeze(1).float() / 255.0)
_, reconstructions = model(Variable(ground_truth).cuda())
reconstruction = reconstructions.cpu().view_as(ground_truth).data
ground_truth_logger.log(
make_grid(ground_truth, nrow=int(BATCH_SIZE ** 0.5), normalize=True, range=(0, 1)).numpy())
reconstruction_logger.log(
make_grid(reconstruction, nrow=int(BATCH_SIZE ** 0.5), normalize=True, range=(0, 1)).numpy())
def test_make_grid_raises_with_variable(self):
t = torch.autograd.Variable(torch.rand(3, 10, 10))
with self.assertRaises(TypeError):
utils.make_grid(t)
with self.assertRaises(TypeError):
utils.make_grid([t, t, t, t])
def visualize_image(self, writer, dataset, image, target, output, global_step):
grid_image = make_grid(image[:3].clone().cpu().data, 3, normalize=True)
writer.add_image('Image', grid_image, global_step)
grid_image = make_grid(decode_seg_map_sequence(torch.max(output[:3], 1)[1].detach().cpu().numpy(),
dataset=dataset), 3, normalize=False, range=(0, 255))
writer.add_image('Predicted label', grid_image, global_step)
grid_image = make_grid(decode_seg_map_sequence(torch.squeeze(target[:3], 1).detach().cpu().numpy(),
dataset=dataset), 3, normalize=False, range=(0, 255))
writer.add_image('Groundtruth label', grid_image, global_step)
def on_step_validation(self, state):
if not self.done:
data = state[self.key].clone()
if len(data.size()) == 3:
data = data.unsqueeze(1)
if self._data is None:
remaining = self.num_images if self.num_images < data.size(0) else data.size(0)
self._data = data[:remaining].to('cpu')
else:
remaining = self.num_images - self._data.size(0)
if remaining > data.size(0):
remaining = data.size(0)
self._data = torch.cat((self._data, data[:remaining].to('cpu')), dim=0)
if self._data.size(0) >= self.num_images:
image = utils.make_grid(
self._data,
nrow=self.nrow,
padding=self.padding,
normalize=self.normalize,
range=self.range,
scale_each=self.scale_each,
pad_value=self.pad_value
)
self._writer.add_image(self.name, image, state[torchbearer.EPOCH])
self.done = True
self._data = None
def _save_originals(dir_z, dir_x, train_test):
dataset = EmbeddingsImagesDataset(dir_z, dir_x)
fixed_dataloader = DataLoader(dataset, 16)
fixed_batch = next(iter(fixed_dataloader))
temp = make_grid(fixed_batch['x'], nrow=4).numpy().transpose((1, 2, 0))
filename_images = os.path.join(self.dir_experiment, 'originals_{}.png'.format(train_test))
Image.fromarray(np.uint8((temp + 1) * 127.5)).save(filename_images)
def _generate_from_model(dir_z, dir_x, train_test):
dataset = EmbeddingsImagesDataset(dir_z, dir_x)
fixed_dataloader = DataLoader(dataset, 16)
fixed_batch = next(iter(fixed_dataloader))
z = Variable(fixed_batch['z']).type(torch.FloatTensor).cuda()
g_z = g.forward(z)
filename_images = os.path.join(self.dir_experiment, 'epoch_{}_{}.png'.format(epoch, train_test))
temp = make_grid(g_z.data[:16], nrow=4).cpu().numpy().transpose((1, 2, 0))
Image.fromarray(np.uint8((temp + 1) * 127.5)).save(filename_images)
def train(self, epoch_to_restore=0):
g = Generator(self.nb_channels_first_layer, self.dim)
if epoch_to_restore > 0:
filename_model = os.path.join(self.dir_models, 'epoch_{}.pth'.format(epoch_to_restore))
g.load_state_dict(torch.load(filename_model))
else:
g.apply(weights_init)
g.cuda()
g.train()
dataset = EmbeddingsImagesDataset(self.dir_z_train, self.dir_x_train)
dataloader = DataLoader(dataset, self.batch_size, shuffle=True, num_workers=4, pin_memory=True)
fixed_dataloader = DataLoader(dataset, 16)
fixed_batch = next(iter(fixed_dataloader))
criterion = torch.nn.L1Loss()
optimizer = optim.Adam(g.parameters())
writer = SummaryWriter(self.dir_logs)
try:
epoch = epoch_to_restore
while True:
g.train()
for _ in range(self.nb_epochs_to_save):
epoch += 1
for idx_batch, current_batch in enumerate(tqdm(dataloader)):
g.zero_grad()
x = Variable(current_batch['x']).type(torch.FloatTensor).cuda()
z = Variable(current_batch['z']).type(torch.FloatTensor).cuda()
g_z = g.forward(z)
loss = criterion(g_z, x)
loss.backward()
optimizer.step()
writer.add_scalar('train_loss', loss, epoch)
z = Variable(fixed_batch['z']).type(torch.FloatTensor).cuda()
g.eval()
g_z = g.forward(z)
images = make_grid(g_z.data[:16], nrow=4, normalize=True)
writer.add_image('generations', images, epoch)
filename = os.path.join(self.dir_models, 'epoch_{}.pth'.format(epoch))
torch.save(g.state_dict(), filename)
finally:
print('[*] Closing Writer.')
writer.close()
def show(self, source, iteration):
if self.cuda:
source = source.cuda()
source = Variable(source)
outputs = self.gen_model(source)
images = vutils.make_grid(outputs.cpu().data, normalize=True)
images = np.transpose(images.numpy(), (1, 2, 0))
plt.imshow(images)
plt.show()
def _generate_random():
nb_samples = 16
z = np.random.randn(nb_samples, self.dim)
norms = np.sqrt(np.sum(z ** 2, axis=1))
norms = np.expand_dims(norms, axis=1)
norms = np.repeat(norms, self.dim, axis=1)
z /= norms
z = Variable(torch.from_numpy(z)).type(torch.FloatTensor).cuda()
g_z = g.forward(z)
filename_images = os.path.join(self.dir_experiment, 'epoch_{}_random.png'.format(epoch))
temp = make_grid(g_z.data[:16], nrow=4).cpu().numpy().transpose((1, 2, 0))
Image.fromarray(np.uint8((temp + 1) * 127.5)).save(filename_images)
def show_landmarks_batch(sample_batched):
"""Show image with landmarks for a batch of samples."""
images_batch, landmarks_batch = \
sample_batched['image'], sample_batched['landmarks']
batch_size = len(images_batch)
im_size = images_batch.size(2)
grid = utils.make_grid(images_batch)
plt.imshow(grid.numpy().transpose((1, 2, 0)))
for i in range(batch_size):
plt.scatter(landmarks_batch[i, :, 0].numpy() + i * im_size,
landmarks_batch[i, :, 1].numpy(),
s=10, marker='.', c='r')
plt.title('Batch from dataloader')
def test_make_grid_not_inplace(self):
t = torch.rand(5, 3, 10, 10)
t_clone = t.clone()
utils.make_grid(t, normalize=False)
assert torch.equal(t, t_clone), 'make_grid modified tensor in-place'
utils.make_grid(t, normalize=True, scale_each=False)
assert torch.equal(t, t_clone), 'make_grid modified tensor in-place'
utils.make_grid(t, normalize=True, scale_each=True)
assert torch.equal(t, t_clone), 'make_grid modified tensor in-place'
def show_grid(input, output):
images = torch.cat((input.data, output.data), 0)
images = make_grid(images, padding=10, normalize=True, scale_each=True).cpu()
image_np = images.numpy().transpose(1, 2, 0)
plt.imshow(image_np)
plt.show()
def train(epoch, phase='train'):
global global_step, best_loss, best_metric, best_accuracy
if phase == 'train':
writer.add_scalar('%s/learning_rate' % phase, get_lr(), epoch)
model.train() if phase == 'train' else model.eval()
torch.set_grad_enabled(True) if phase == 'train' else torch.set_grad_enabled(False)
dataloader = train_dataloader if phase == 'train' else valid_dataloader
running_loss, running_metric, running_accuracy = 0.0, 0.0, 0.0
worst_loss, worst_metric = best_loss, best_metric
it, total = 0, 0
if phase == 'valid':
total_probs = []
total_truth = []
pbar_disable = False if epoch == start_epoch else None
pbar = tqdm(dataloader, unit="images", unit_scale=dataloader.batch_size, disable=pbar_disable)
for batch in pbar:
image_ids, inputs, targets = batch['image_id'], batch['input'], batch['mask']
if use_gpu:
inputs = inputs.cuda()
targets = targets.cuda()
# forward
logit, logit_pixel, logit_image = model(inputs)
# look at the center only
if args.loss_on_center:
logit = remove_padding(logit)
logit_pixel = (remove_padding(l) for l in logit_pixel)
targets = remove_padding(targets)
truth_pixel = targets
truth_image = (truth_pixel.sum(dim=(1, 2)) > 0).float()
loss = models.deep_supervised_criterion(logit, logit_pixel, logit_image, truth_pixel, truth_image)
if not args.loss_on_center and not args.resize:
logit = remove_padding(logit)
targets = remove_padding(targets)
probs = torch.sigmoid(logit).squeeze(1)
# predictions = probs.squeeze(1) > 0.5
if phase == 'train':
# backward
optimizer.backward(loss / args.gradient_accumulation)
if it % args.gradient_accumulation == 0:
optimizer.step()
optimizer.zero_grad()
# statistics
it += 1
global_step += 1
loss = loss.item()
running_loss += (loss * targets.size(0))
total += targets.size(0)
writer.add_scalar('%s/loss' % phase, loss, global_step)
targets_numpy = targets.cpu().numpy()
probs_numpy = probs.cpu().detach().numpy()
predictions_numpy = probs_numpy > 0.5 # predictions.cpu().numpy()
metric_array = calc_metric(targets_numpy, predictions_numpy, type='iou', size_average=False)
metric = metric_array.mean()
running_metric += metric_array.sum()
running_accuracy += calc_metric(targets_numpy, predictions_numpy, type='pixel_accuracy',
size_average=False).sum()
if phase == 'valid':
total_truth.append(targets_numpy)
total_probs.append(probs_numpy)
visualize_output = False
if worst_loss > loss:
worst_loss = loss
visualize_output = True
if worst_metric < metric:
worst_metric = metric
visualize_output = True
if visualize_output and args.debug:
# sort samples by metric
ind = np.argsort(metric_array)
images = remove_padding(inputs.cpu())
images = images[ind]
probs = probs[ind].cpu()
predictions = predictions[ind].cpu()
targets = targets[ind].cpu()
preds = torch.cat([probs] * 3, 1)
mask = torch.cat([targets.unsqueeze(1)] * 3, 1)
all = images.clone()
all[:, 0] = torch.max(images[:, 0], predictions.float())
all[:, 1] = torch.max(images[:, 1], targets)
all = torch.cat((torch.cat((all, images), 3), torch.cat((preds, mask), 3)), 2)
all_grid = vutils.make_grid(all, nrow=4, normalize=False, pad_value=1)
writer.add_image('%s/img-mask-pred' % phase, all_grid, global_step)
# update the progress bar
pbar.set_postfix({
'loss': "%.05f" % (running_loss / total),
'metric': "%.03f" % (running_metric / total)
})
epoch_loss = running_loss / total
epoch_metric = running_metric / total
epoch_accuracy = running_accuracy / total
writer.add_scalar('%s/metric' % phase, epoch_metric, epoch)
writer.add_scalar('%s/accuracy' % phase, epoch_accuracy, epoch)
writer.add_scalar('%s/epoch_loss' % phase, epoch_loss, epoch)
if phase == 'valid':
def save_checkpoint(name):
cycle = ('-cycle%d' % (epoch // args.lr_scheduler_step_size)) if args.lr_scheduler == 'clr' else ''
model_name = name + '-model'
model_file_name = '%d-%s-%s%s.pth' % (start_timestamp, model_name, full_name, cycle)
model_file = models_dir / model_file_name
models.save(model, model_file)
mode_file_simple = Path(models_dir / (model_name + '-%s%s.pth' % (args.data_fold, cycle)))
if mode_file_simple.is_symlink() or mode_file_simple.exists():
mode_file_simple.unlink()
mode_file_simple.symlink_to(model_file.relative_to(mode_file_simple.parent))
checkpoint = {
'epoch': epoch,
'step': global_step,
'model_file': str(model_file),
'best_loss': best_loss,
'best_metric': best_metric,
'best_accuracy': best_accuracy,
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict()
}
checkpoint_filename = name + '-checkpoint-%s%s.pth' % (full_name, cycle)
checkpoint_file = checkpoint_dir / checkpoint_filename
torch.save(checkpoint, checkpoint_file)
checkpoint_file_simple = Path(checkpoint_dir / (name + '-checkpoint-%s%s.pth' % (args.data_fold, cycle)))
if checkpoint_file_simple.is_symlink() or checkpoint_file_simple.exists():
checkpoint_file_simple.unlink()
checkpoint_file_simple.symlink_to(checkpoint_file.relative_to(checkpoint_file_simple.parent))
if epoch_loss < best_loss:
best_loss = epoch_loss
save_checkpoint('best-loss')
if epoch_metric > best_metric:
best_metric = epoch_metric
save_checkpoint('best-metric')
if epoch_accuracy > best_accuracy:
best_accuracy = epoch_accuracy
save_checkpoint('best-accuracy')
save_checkpoint('last')
return epoch_loss, epoch_metric, epoch_accuracy
def train(self):
self.optimizer_G = optim.Adam(self.G.parameters(), self.config.g_lr, [self.config.beta1, self.config.beta2])
self.optimizer_D = optim.Adam(self.D.parameters(), self.config.d_lr, [self.config.beta1, self.config.beta2])
self.lr_scheduler_G = optim.lr_scheduler.StepLR(self.optimizer_G, step_size=self.config.lr_decay_iters, gamma=0.1)
self.lr_scheduler_D = optim.lr_scheduler.StepLR(self.optimizer_D, step_size=self.config.lr_decay_iters, gamma=0.1)
self.load_checkpoint()
if self.cuda and self.config.ngpu > 1:
self.G = nn.DataParallel(self.G, device_ids=list(range(self.config.ngpu)))
self.D = nn.DataParallel(self.D, device_ids=list(range(self.config.ngpu)))
val_iter = iter(self.data_loader.val_loader)
x_sample, c_org_sample = next(val_iter)
x_sample = x_sample.to(self.device)
c_sample_list = self.create_labels(c_org_sample, self.config.attrs)
c_sample_list.insert(0, c_org_sample) # reconstruction
self.g_lr = self.lr_scheduler_G.get_lr()[0]
self.d_lr = self.lr_scheduler_D.get_lr()[0]
data_iter = iter(self.data_loader.train_loader)
start_time = time.time()
for i in range(self.current_iteration, self.config.max_iters):
self.G.train()
self.D.train()
# =================================================================================== #
# 1. Preprocess input data #
# =================================================================================== #
# fetch real images and labels
try:
x_real, label_org = next(data_iter)
except:
data_iter = iter(self.data_loader.train_loader)
x_real, label_org = next(data_iter)
# generate target domain labels randomly
rand_idx = torch.randperm(label_org.size(0))
label_trg = label_org[rand_idx]
c_org = label_org.clone()
c_trg = label_trg.clone()
x_real = x_real.to(self.device) # input images
c_org = c_org.to(self.device) # original domain labels
c_trg = c_trg.to(self.device) # target domain labels
label_org = label_org.to(self.device) # labels for computing classification loss
label_trg = label_trg.to(self.device) # labels for computing classification loss
# =================================================================================== #
# 2. Train the discriminator #
# =================================================================================== #
# compute loss with real images
out_src, out_cls = self.D(x_real)
d_loss_real = - torch.mean(out_src)
d_loss_cls = self.classification_loss(out_cls, label_org)
# compute loss with fake images
attr_diff = c_trg - c_org
attr_diff = attr_diff * torch.rand_like(attr_diff) * (2 * self.config.thres_int)
x_fake = self.G(x_real, attr_diff)
out_src, out_cls = self.D(x_fake.detach())
d_loss_fake = torch.mean(out_src)
# compute loss for gradient penalty
alpha = torch.rand(x_real.size(0), 1, 1, 1).to(self.device)
x_hat = (alpha * x_real.data + (1 - alpha) * x_fake.data).requires_grad_(True)
out_src, _ = self.D(x_hat)
d_loss_gp = self.gradient_penalty(out_src, x_hat)
# backward and optimize
d_loss_adv = d_loss_real + d_loss_fake + self.config.lambda_gp * d_loss_gp
d_loss = d_loss_adv + self.config.lambda1 * d_loss_cls
self.optimizer_D.zero_grad()
d_loss.backward(retain_graph=True)
self.optimizer_D.step()
# summarize
scalars = {}
scalars['D/loss'] = d_loss.item()
scalars['D/loss_adv'] = d_loss_adv.item()
scalars['D/loss_cls'] = d_loss_cls.item()
scalars['D/loss_real'] = d_loss_real.item()
scalars['D/loss_fake'] = d_loss_fake.item()
scalars['D/loss_gp'] = d_loss_gp.item()
# =================================================================================== #
# 3. Train the generator #
# =================================================================================== #
if (i + 1) % self.config.n_critic == 0:
# original-to-target domain
x_fake = self.G(x_real, attr_diff)
out_src, out_cls = self.D(x_fake)
g_loss_adv = - torch.mean(out_src)
g_loss_cls = self.classification_loss(out_cls, label_trg)
# target-to-original domain
x_reconst = self.G(x_fake, c_org - c_org)
g_loss_rec = torch.mean(torch.abs(x_real - x_reconst))
# backward and optimize
g_loss = g_loss_adv + self.config.lambda3 * g_loss_rec + self.config.lambda2 * g_loss_cls
self.optimizer_G.zero_grad()
g_loss.backward()
self.optimizer_G.step()
# summarize
scalars['G/loss'] = g_loss.item()
scalars['G/loss_adv'] = g_loss_adv.item()
scalars['G/loss_cls'] = g_loss_cls.item()
scalars['G/loss_rec'] = g_loss_rec.item()
self.current_iteration += 1
# =================================================================================== #
# 4. Miscellaneous #
# =================================================================================== #
if self.current_iteration % self.config.summary_step == 0:
et = time.time() - start_time
et = str(datetime.timedelta(seconds=et))[:-7]
print('Elapsed [{}], Iteration [{}/{}]'.format(et, self.current_iteration, self.config.max_iters))
for tag, value in scalars.items():
self.writer.add_scalar(tag, value, self.current_iteration)
if self.current_iteration % self.config.sample_step == 0:
self.G.eval()
with torch.no_grad():
x_sample = x_sample.to(self.device)
x_fake_list = [x_sample]
for c_trg_sample in c_sample_list:
attr_diff = c_trg_sample.to(self.device) - c_org_sample.to(self.device)
attr_diff = attr_diff * self.config.thres_int
x_fake_list.append(self.G(x_sample, attr_diff.to(self.device)))
x_concat = torch.cat(x_fake_list, dim=3)
self.writer.add_image('sample', make_grid(self.denorm(x_concat.data.cpu()), nrow=1),
self.current_iteration)
save_image(self.denorm(x_concat.data.cpu()),
os.path.join(self.config.sample_dir, 'sample_{}.jpg'.format(self.current_iteration)),
nrow=1, padding=0)
if self.current_iteration % self.config.checkpoint_step == 0:
self.save_checkpoint()
self.lr_scheduler_G.step()
self.lr_scheduler_D.step()
batch_v = batch_v.to(device)
gen_output_v = net_gener(gen_input_v)
# train discriminator
dis_optimizer.zero_grad()
dis_output_true_v = net_discr(batch_v)
dis_output_fake_v = net_discr(gen_output_v.detach())
dis_loss = objective(dis_output_true_v, true_labels_v) + objective(dis_output_fake_v, fake_labels_v)
dis_loss.backward()
dis_optimizer.step()
dis_losses.append(dis_loss.item())
# train generator
gen_optimizer.zero_grad()
dis_output_v = net_discr(gen_output_v)
gen_loss_v = objective(dis_output_v, true_labels_v)
gen_loss_v.backward()
gen_optimizer.step()
gen_losses.append(gen_loss_v.item())
iter_no += 1
if iter_no % REPORT_EVERY_ITER == 0:
log.info("Iter %d: gen_loss=%.3e, dis_loss=%.3e", iter_no, np.mean(gen_losses), np.mean(dis_losses))
writer.add_scalar("gen_loss", np.mean(gen_losses), iter_no)
writer.add_scalar("dis_loss", np.mean(dis_losses), iter_no)
gen_losses = []
dis_losses = []
if iter_no % SAVE_IMAGE_EVERY_ITER == 0:
writer.add_image("fake", vutils.make_grid(gen_output_v.data[:64]), iter_no)
writer.add_image("real", vutils.make_grid(batch_v.data[:64]), iter_no)
def __write_images(image_outputs, display_image_num, file_name):
image_outputs = [images.expand(-1, 3, -1, -1) for images in image_outputs] # expand gray-scale images to 3 channels
image_tensor = torch.cat([images[:display_image_num] for images in image_outputs], 0)
image_grid = vutils.make_grid(image_tensor.data, nrow=display_image_num, padding=0, normalize=True)
vutils.save_image(image_grid, file_name, nrow=1)
input_a16 = align_to_num(input_cv2, 16)
input_a16 = to_tensor(input_a16, device)
# target
target_cv2 = cv2.imread(os.path.join('testsets/%s/gt' % testset, img.replace('rain', 'clean')))
target = align_to_num(target_cv2)
target = to_tensor(target, device)
# target align to 16
target_a16 = align_to_num(target_cv2, 16)
target_a16 = to_tensor(target_a16, device)
# initial for measurement
cal_input = input
cal_target = target
start_time = time.time()
# rainmap
est = estNet(input)
logimg = make_grid(est.data.clamp(0., 1.), nrow=8, normalize=False, scale_each=False)
est = logimg.mul_(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to('cpu', torch.uint8).numpy()[:, :, ::-1]
derain = align_to_num(est)
rainmap = make_rainmap(input_cv2, derain)
rainmap = to_tensor(rainmap, device)
# edge
derain = prepare_image_cv2(np.array(est, dtype=np.float32))
derain_in = derain.transpose((1, 2, 0))
scale = [0.5, 1, 1.5]
_, H, W = derain.shape
multi_fuse = np.zeros((H, W), np.float32)
for k in range(0, len(scale)):
im_ = cv2.resize(derain_in, None, fx=scale[k], fy=scale[k], interpolation=cv2.INTER_LINEAR)
im_ = im_.transpose((2, 0, 1))
edges = rcfNet(torch.unsqueeze(torch.from_numpy(im_).to(device), 0))
edge = torch.squeeze(edges[-1].detach()).cpu().numpy()
# Backward the averaged gradient
loss /= p['nAveGrad']
loss.backward()
aveGrad += 1
# Update the weights once in p['nAveGrad'] forward passes
if aveGrad % p['nAveGrad'] == 0:
writer.add_scalar('data/total_loss_iter', loss.item(), ii + num_img_tr * epoch)
optimizer.step()
optimizer.zero_grad()
aveGrad = 0
if ii % num_img_tr / 20 == 0:
grid_image = make_grid(inputs[:3].clone().cpu().data, 3, normalize=True)
writer.add_image('image', grid_image)
grid_image = make_grid(utils.decode_seg_map_sequence(torch.max(output[:3], 1)[1].detach().cpu().numpy()), 3, normalize=False,
range=(0, 255))
writer.add_image('Predicted label', grid_image)
grid_image = make_grid(utils.decode_seg_map_sequence(torch.squeeze(gts[:3], 1).detach().cpu().numpy()), 3, normalize=False, range=(0, 255))
writer.add_image('Groundtruth label', grid_image)
# Save the model
if (epoch % snapshot) == snapshot - 1:
torch.save(net.state_dict(), os.path.join(save_dir, 'models', modelName + '_epoch-' + str(epoch) + '.pth'))
print("Save model at {}\n".format(os.path.join(save_dir, 'models', modelName + '_epoch-' + str(epoch) + '.pth')))
# One testing epoch
if useTest and epoch % nTestInterval == (nTestInterval - 1):
net.eval()
train_loss.append(loss.data[0])
G_loss.backward()
Q_solver.step()
reset_grad()
print('D-loss: {:.4f}, G-loss: {:.4f} time: {:.4f} seconds'.format(D_loss.data.cpu().numpy(),G_loss.data.cpu().numpy(), time()-tic))
if epoch % 10 == 0:
torch.save(model,"model_latest.pth")
if args.cuda:
X_real, X_fake = X_fake.cpu(), X_real.cpu()
plt.subplot(321)
grid = make_grid(X_real.data, nrow=4).numpy()
grid = np.moveaxis(grid,0,-1)
plt.imshow(cv2.cvtColor(grid, cv2.COLOR_BGR2RGB))
plt.subplot(322)
grid = make_grid(X_fake.data, nrow=4).numpy()
grid = np.moveaxis(grid,0,-1)
plt.imshow(cv2.cvtColor(grid, cv2.COLOR_BGR2RGB))
#plt.show()
out_dir = "output"
if not os.path.exists(out_dir):
os.makedirs(out_dir)
plt.savefig('{}/{}.png'.format(out_dir,str(epoch)), bbox_inches='tight')
torch.save(model,'{}/checkpoint_{}.pth.tar'.format(out_dir,str(epoch)))
def train_disc_gen(self, dataloader, nz, netD, netG, optimizerD,
optimizerG, num_epochs, device):
# Lists to keep track of progress
img_list = []
G_losses = []
D_losses = []
iters = 0
# Establish convention for real and fake labels during training
real_label = 1.
fake_label = 0.
criterion = nn.BCELoss() # BCELoss function
# Create batch of latent vectors that we will use to visualize
# the progression of the generator
fixed_noise = torch.randn(64, nz, 1, 1, device=device)
print("Starting Training Loop...")
# For each epoch
for epoch in range(num_epochs):
# For each batch in the dataloader
for i, data in enumerate(dataloader, 0):
############################
# (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
###########################
## Train with all-real batch
netD.zero_grad()
# Format batch
real_cpu = data[0].to(device)
b_size = real_cpu.size(0)
label = torch.full((b_size, ),
real_label,
dtype=torch.float,
device=device)
# Forward pass real batch through D
output = netD(real_cpu).view(-1)
# Calculate loss on all-real batch
errD_real = criterion(output, label)
# Calculate gradients for D in backward pass
errD_real.backward()
D_x = output.mean().item()
## Train with all-fake batch
# Generate batch of latent vectors
noise = torch.randn(b_size, nz, 1, 1, device=device)
# Generate fake image batch with G
fake = netG(noise)
label.fill_(fake_label)
# Classify all fake batch with D
output = netD(fake.detach()).view(-1)
# Calculate D's loss on the all-fake batch
errD_fake = criterion(output, label)
# Calculate the gradients for this batch
errD_fake.backward()
D_G_z1 = output.mean().item()
# Add the gradients from the all-real and all-fake batches
errD = errD_real + errD_fake
# Update D
optimizerD.step()
############################
# (2) Update G network: maximize log(D(G(z)))
###########################
netG.zero_grad()
label.fill_(
real_label) # fake labels are real for generator cost
# Since we just updated D, perform another forward pass of all-fake batch through D
output = netD(fake).view(-1)
# Calculate G's loss based on this output
errG = criterion(output, label)
# Calculate gradients for G
errG.backward()
D_G_z2 = output.mean().item()
# Update G
optimizerG.step()
# Output training stats
if i % 50 == 0:
print(
'[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
% (epoch, num_epochs, i, len(dataloader), errD.item(),
errG.item(), D_x, D_G_z1, D_G_z2))
# Save Losses for plotting later
G_losses.append(errG.item())
D_losses.append(errD.item())
# Check how the generator is doing by saving G's output on fixed_noise
if (iters % 500 == 0) or ((epoch == num_epochs - 1) and
(i == len(dataloader) - 1)):
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
img_list.append(
vutils.make_grid(fake, padding=2, normalize=True))
iters += 1
#---------------------end training loop-----------------
return img_list, G_losses, D_losses
running_loss_tr = running_loss_tr / 10 / crop_size
print('[Epoch:%d [%d/%d], numImages: %5d]' %
(epoch, ii, num_img_tr,
ii * p['trainBatch'] + inputs.data.shape[0]))
print('(poly lr policy) learning rate: ', lr_)
print('Loss: %f' % running_loss_tr)
running_loss_tr = 0
stop_time = timeit.default_timer()
print("Execution time: " + str(stop_time - start_time) + "\n")
start_time = timeit.default_timer()
# Show 10 * 3 images results each epoch
# 每轮epoch输出可视化结果
if ii % num_img_tr == (num_img_tr - 1):
grid_image = make_grid(inputs[:3].clone().cpu().data,
3,
normalize=True)
writer.add_image('Image', grid_image, global_step)
grid_image = make_grid(utils.decode_seg_map_sequence(
torch.max(outputs[:3], 1)[1].detach().cpu().numpy()),
3,
normalize=False,
range=(0, 255))
writer.add_image('Predicted label', grid_image, global_step)
grid_image = make_grid(utils.decode_seg_map_sequence(
torch.squeeze(labels[:3], 1).detach().cpu().numpy()),
3,
normalize=False,
range=(0, 255))
writer.add_image('Groundtruth label', grid_image, global_step)
def get_labelname(self):
return self.label_names
if __name__ == "__main__":
data_transform = transforms.Compose(
[transforms.ToTensor(), ])
train_datasets = CharacterDateset('../data', type='train')
train_dataloader = DataLoader(train_datasets, batch_size=100, shuffle=True, num_workers=2)
val_datasets = CharacterDateset('../data', type='val')
val_dataloader = DataLoader(val_datasets, batch_size=4, shuffle=True, num_workers=2)
# print(val_datasets.get_img_path_list())
# train_datasets.save_label_map()
# for i, (img, label) in enumerate(val_dataloader):
# print(f"{i}:{img}")
labels_maps = val_datasets.get_labelname()
dataiter = iter(val_dataloader)
images, labels = dataiter.next()
print(images.shape)
imshow(make_grid(images))
print(' '.join('%5s' % labels_maps[labels[j]] for j in range(4)))
print(len(val_dataloader))
im = Image.open("../data/test1_/5bf00c19cc06f66070a6aad27336b62e40f14fab.jpg").convert("L").resize((256, 256))
# im.show(title="origin")
im_1 = RandomColor(im)
im_2 = RandomGaussian(im_1, mean=0.2, sigma=0.3)
# im_2.show(title="convert")
print(val_datasets.get_labelname())
def main(args):
# parameters
img_size = 128
z_dim = 128
lamb_obj = 1.0
lamb_app = 1.0
lamb_img = 0.1
num_classes = 184 if args.dataset == 'coco' else 179
num_obj = 8 if args.dataset == 'coco' else 31
args.out_path = os.path.join(args.out_path, args.dataset + '_1gpu', str(img_size))
# data loader
train_data = get_dataset(args.dataset, img_size)
num_gpus = torch.cuda.device_count()
num_workers = 2
if num_gpus > 1:
parallel = True
args.batch_size = args.batch_size * num_gpus
num_workers = num_workers * num_gpus
else:
parallel = False
print("{} GPUs, {} workers are used".format(num_gpus, num_workers))
dataloader = torch.utils.data.DataLoader(
train_data, batch_size=args.batch_size,
drop_last=True, shuffle=True, num_workers=num_workers)
# Load model
device = torch.device('cuda')
netG = ResnetGenerator128(num_classes=num_classes, output_dim=3).to(device)
netD = CombineDiscriminator128_app(num_classes=num_classes).to(device)
# if os.path.isfile(args.checkpoint):
# state_dict = torch.load(args.checkpoint)
# parallel = True
if parallel:
netG = DataParallelWithCallback(netG)
netD = nn.DataParallel(netD)
g_lr, d_lr = args.g_lr, args.d_lr
gen_parameters = []
for key, value in dict(netG.named_parameters()).items():
if value.requires_grad:
if 'mapping' in key:
gen_parameters += [{'params': [value], 'lr': g_lr * 0.1}]
else:
gen_parameters += [{'params': [value], 'lr': g_lr}]
g_optimizer = torch.optim.Adam(gen_parameters, betas=(0, 0.999))
dis_parameters = []
for key, value in dict(netD.named_parameters()).items():
if value.requires_grad:
dis_parameters += [{'params': [value], 'lr': d_lr}]
d_optimizer = torch.optim.Adam(dis_parameters, betas=(0, 0.999))
# make dirs
if not os.path.exists(args.out_path):
os.makedirs(args.out_path)
if not os.path.exists(os.path.join(args.out_path, 'model/')):
os.makedirs(os.path.join(args.out_path, 'model/'))
writer = SummaryWriter(os.path.join(args.out_path, 'log'))
# writer = None
logger = setup_logger("lostGAN", args.out_path, 0)
logger.info(netG)
logger.info(netD)
start_time = time.time()
vgg_loss = VGGLoss()
vgg_loss = nn.DataParallel(vgg_loss)
l1_loss = nn.DataParallel(nn.L1Loss())
for epoch in range(args.total_epoch):
netG.train()
netD.train()
print("Epoch {}/{}".format(epoch, args.total_epoch))
for idx, data in enumerate(tqdm(dataloader)):
real_images, label, bbox = data
real_images, label, bbox = real_images.to(device), label.long().to(device).unsqueeze(-1), bbox.float()
# update D network
netD.zero_grad()
real_images, label = real_images.to(device), label.long().to(device)
d_out_real, d_out_robj, d_out_robj_app = netD(real_images, bbox, label)
d_loss_real = torch.nn.ReLU()(1.0 - d_out_real).mean()
d_loss_robj = torch.nn.ReLU()(1.0 - d_out_robj).mean()
d_loss_robj_app = torch.nn.ReLU()(1.0 - d_out_robj_app).mean()
# print(d_loss_robj)
# print(d_loss_robj_app)
z = torch.randn(real_images.size(0), num_obj, z_dim).to(device)
fake_images = netG(z, bbox, y=label.squeeze(dim=-1))
d_out_fake, d_out_fobj, d_out_fobj_app = netD(fake_images.detach(), bbox, label)
d_loss_fake = torch.nn.ReLU()(1.0 + d_out_fake).mean()
d_loss_fobj = torch.nn.ReLU()(1.0 + d_out_fobj).mean()
d_loss_fobj_app = torch.nn.ReLU()(1.0 + d_out_fobj_app).mean()
d_loss = lamb_obj * (d_loss_robj + d_loss_fobj) + lamb_img * (d_loss_real + d_loss_fake) + lamb_app * (d_loss_robj_app + d_loss_fobj_app)
d_loss.backward()
d_optimizer.step()
# update G network
if (idx % 1) == 0:
netG.zero_grad()
g_out_fake, g_out_obj, g_out_obj_app = netD(fake_images, bbox, label)
g_loss_fake = - g_out_fake.mean()
g_loss_obj = - g_out_obj.mean()
g_loss_obj_app = - g_out_obj_app.mean()
pixel_loss = l1_loss(fake_images, real_images).mean()
feat_loss = vgg_loss(fake_images, real_images).mean()
g_loss = g_loss_obj * lamb_obj + g_loss_fake * lamb_img + pixel_loss + feat_loss + lamb_app * g_loss_obj_app
g_loss.backward()
g_optimizer.step()
# print("d_loss_real={:.3f}, d_loss_robj={:.3f}, d_loss_robj_app={:.3f}".format(d_loss_real.item(), d_loss_robj.item(), d_loss_robj_app.item()))
# print("d_loss_fake={:.3f}, d_loss_fobj={:.3f}, d_loss_fobj_app={:.3f}".format(d_loss_fake.item(), d_loss_fobj.item(), d_loss_fobj_app.item()))
# print("g_loss_fake={:.3f}, g_loss_obj={:.3f}, g_loss_obj_app={:.3f}".format(g_loss_fake.item(), g_loss_obj.item(), g_loss_obj_app.item()))
if (idx + 1) % 100 == 0:
elapsed = time.time() - start_time
elapsed = str(datetime.timedelta(seconds=elapsed))
logger.info("Time Elapsed: [{}]".format(elapsed))
logger.info("Step[{}/{}], d_out_real: {:.4f}, d_out_fake: {:.4f}, g_out_fake: {:.4f} ".format(epoch + 1,
idx + 1,
d_loss_real.item(),
d_loss_fake.item(),
g_loss_fake.item()))
logger.info(" d_obj_real: {:.4f}, d_obj_fake: {:.4f}, g_obj_fake: {:.4f} ".format(
d_loss_robj.item(),
d_loss_fobj.item(),
g_loss_obj.item()))
logger.info(" d_obj_real_app: {:.4f}, d_obj_fake_app: {:.4f}, g_obj_fake_app: {:.4f} ".format(
d_loss_robj_app.item(),
d_loss_fobj_app.item(),
g_loss_obj_app.item()))
logger.info(" pixel_loss: {:.4f}, feat_loss: {:.4f}".format(pixel_loss.item(), feat_loss.item()))
if writer is not None:
writer.add_image("real images", make_grid(real_images.cpu().data * 0.5 + 0.5, nrow=4), epoch * len(dataloader) + idx + 1)
writer.add_image("fake images", make_grid(fake_images.cpu().data * 0.5 + 0.5, nrow=4), epoch * len(dataloader) + idx + 1)
writer.add_scalars("D_loss_real", {"real": d_loss_real.item(),
"robj": d_loss_robj.item(),
"robj_app": d_loss_robj_app.item(),
"loss": d_loss.item()})
writer.add_scalars("D_loss_fake", {"fake": d_loss_fake.item(),
"fobj": d_loss_fobj.item(),
"fobj_app": d_loss_fobj_app.item()})
writer.add_scalars("G_loss", {"fake": g_loss_fake.item(),
"obj_app": g_loss_obj_app.item(),
"obj": g_loss_obj.item(),
"loss": g_loss.item()})
# save model
if (epoch + 1) % 5 == 0:
torch.save(netG.state_dict(), os.path.join(args.out_path, 'model/', 'G_%d.pth' % (epoch + 1)))
torch.save(netD.state_dict(), os.path.join(args.out_path, 'model/', 'D_%d.pth' % (epoch + 1)))
with torch.no_grad():
fake = netG(noise).detach()
vutils.save_image(fake.data,
'%s/fake_img_epoch_%03d.jpg' %
("./celeba/load_weight_progress", epoch + 1),
normalize=True)
with torch.no_grad():
fake = netG(fixed_noise).detach()
vutils.save_image(fake.data,
'%s/fake_img_version_epoch_%03d.jpg' %
("./celeba/load_weight_results2", epoch + 1),
normalize=True)
img_list.append(
vutils.make_grid(fake, padding=2, normalize=True))
plt.figure(figsize=(10, 5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses, label="G")
plt.plot(D_losses, label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()
real_batch = next(iter(dataloader))
plt.figure(figsize=(24, 24))
plt.subplot(1, 2, 1)
plt.axis("off")
plt.title("Real Images")
i += 1
if opt.cuda:
real_cim, real_vim, real_sim = real_cim.cuda(), real_vim.cuda(
), real_sim.cuda()
mask = torch.cat([
torch.rand(1, 1, maskS, maskS).ge(X.rvs(1)[0]).float()
for _ in range(opt.batchSize)
], 0).cuda()
hint = torch.cat((real_vim * mask, mask), 1)
if flag: # fix samples
writer.add_image(
'target imgs',
vutils.make_grid(real_cim.mul(0.5).add(0.5), nrow=16))
writer.add_image(
'sketch imgs',
vutils.make_grid(real_sim.mul(0.5).add(0.5), nrow=16))
writer.add_image(
'hint',
vutils.make_grid((real_vim * mask).mul(0.5).add(0.5),
nrow=16))
vutils.save_image(
real_cim.mul(0.5).add(0.5),
'%s/color_samples' % opt.outf + '.png')
vutils.save_image(
real_sim.mul(0.5).add(0.5),
'%s/blur_samples' % opt.outf + '.png')
fixed_sketch.resize_as_(real_sim).copy_(real_sim)
fixed_hint.resize_as_(hint).copy_(hint)
dataloader = DataLoader(dataset, batch_size=b_size,
shuffle=True, num_workers=workers)
# Important! The following line sets the device.
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
import matplotlib.pyplot as plt
import numpy as np
import torchvision.utils as vision_utils
real_batch = next(iter(dataloader))
plt.figure(figsize=(8,8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(vision_utils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
import torch.nn as nn
# Will be called by the Generator and Discriminator networks
def weights_init(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(m.weight.data, 0., .02)
elif classname.find('BatchNorm') != -1:
nn.init.normal_(m.weight.data, 1., .02)
nn.init.constant_(m.bias.data, 0)
class Generator(nn.Module):
def __init__(self, ngpu):
super(Generator, self).__init__()
def validate_on_test(self, step, summ_writer=None):
# first put the model in eval mode
self.model.eval()
# everytime make the val list empty
#val_results_list = list()
#cluster_id_list = list()
print("start validate on test")
assert not self.model.training, "Model should be in eval mode"
if self.val_data_nn is None:
test_data_path = self.config.data_paths['test']
# now form the data-loader with the valset path
val_dataloader = torch.utils.data.DataLoader(MuJoCoOfflineData(
config=self.config,
dataset_path=test_data_path,
plot=False,
train=False,
fix_view=True,
num_workers=1),
batch_size=2,
shuffle=True,
drop_last=True)
print(f'Length of val_data is {len(val_dataloader)}')
self.val_data_nn = self.get_data(val_dataloader, summ_writer, step)
print("finish loading data")
val_results_list, cluster_id_list, success_rate_list, feed = self.get_features_from_data_list(
self.val_data_nn, step)
#val_results_list, cluster_id_list, feed = self.get_features(val_dataloader, summ_writer, step)
# now that you have results think about how can you collate and do nearest neighbor
vis_nearest_neighbours, recall_1 = self.compute_nearest_neighbours_dp(
val_results_list, cluster_id_list)
# now the only thing that remains is plotting this on tensorboard
# just to satisfy my paranoia I will also save the matplotlib images
# but after 500 iterations
if feed['global_step'] % 5000 == 0:
n_rows = len(vis_nearest_neighbours)
n_cols = len(vis_nearest_neighbours[0])
fig_size = 2 * np.asarray([n_rows, n_cols])
fig, axes = plt.subplots(nrows=n_rows,
ncols=n_cols,
figsize=fig_size,
sharex=True,
sharey=True)
for i in range(n_rows):
for j in range(n_cols):
axes[i][j].imshow(vis_nearest_neighbours[i][j])
# save the figure and you are done
fig.savefig(f'{self.nn_results_dir}/nn_result_{step}.jpg')
plt.close()
H, W, C = list(vis_nearest_neighbours[0][1].shape)
# finally add it to tensorboard and you are done !!!
# #test_data x (#top_retrieve + 1)
vis_nearest_neighbours = np.stack(vis_nearest_neighbours, axis=0)
# convert to torch
vis_nearest = torch.from_numpy(vis_nearest_neighbours).permute(
0, 1, 4, 2, 3)
# resize
vis_nearest = vis_nearest.view(-1, C, H, W)
# make the grid
grid = make_grid(vis_nearest, nrow=11)
# add it to the tensorboard
feed['writer'].add_scalar('val_nn_recall@1', recall_1, step)
feed['writer'].add_image('val_nn/imgs', grid, step)
def test_tensor2img():
tensor_4d_1 = torch.FloatTensor(2, 3, 4, 4).uniform_(0, 1)
tensor_4d_2 = torch.FloatTensor(1, 3, 4, 4).uniform_(0, 1)
tensor_4d_3 = torch.FloatTensor(3, 1, 4, 4).uniform_(0, 1)
tensor_4d_4 = torch.FloatTensor(1, 1, 4, 4).uniform_(0, 1)
tensor_3d_1 = torch.FloatTensor(3, 4, 4).uniform_(0, 1)
tensor_3d_2 = torch.FloatTensor(3, 6, 6).uniform_(0, 1)
tensor_3d_3 = torch.FloatTensor(1, 6, 6).uniform_(0, 1)
tensor_2d = torch.FloatTensor(4, 4).uniform_(0, 1)
with pytest.raises(TypeError):
# input is not a tensor
tensor2img(4)
with pytest.raises(TypeError):
# input is not a list of tensors
tensor2img([tensor_3d_1, 4])
with pytest.raises(ValueError):
# unsupported 5D tensor
tensor2img(torch.FloatTensor(2, 2, 3, 4, 4).uniform_(0, 1))
# 4d
rlt = tensor2img(tensor_4d_1, out_type=np.uint8, min_max=(0, 1))
tensor_4d_1_np = make_grid(tensor_4d_1, nrow=1, normalize=False).numpy()
tensor_4d_1_np = np.transpose(tensor_4d_1_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt, (tensor_4d_1_np * 255).round())
rlt = tensor2img(tensor_4d_2, out_type=np.uint8, min_max=(0, 1))
tensor_4d_2_np = tensor_4d_2.squeeze().numpy()
tensor_4d_2_np = np.transpose(tensor_4d_2_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt, (tensor_4d_2_np * 255).round())
rlt = tensor2img(tensor_4d_3, out_type=np.uint8, min_max=(0, 1))
tensor_4d_3_np = make_grid(tensor_4d_3, nrow=1, normalize=False).numpy()
tensor_4d_3_np = np.transpose(tensor_4d_3_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt, (tensor_4d_3_np * 255).round())
rlt = tensor2img(tensor_4d_4, out_type=np.uint8, min_max=(0, 1))
tensor_4d_4_np = tensor_4d_4.squeeze().numpy()
np.testing.assert_almost_equal(rlt, (tensor_4d_4_np * 255).round())
# 3d
rlt = tensor2img([tensor_3d_1, tensor_3d_2],
out_type=np.uint8,
min_max=(0, 1))
tensor_3d_1_np = tensor_3d_1.numpy()
tensor_3d_1_np = np.transpose(tensor_3d_1_np[[2, 1, 0], :, :], (1, 2, 0))
tensor_3d_2_np = tensor_3d_2.numpy()
tensor_3d_2_np = np.transpose(tensor_3d_2_np[[2, 1, 0], :, :], (1, 2, 0))
np.testing.assert_almost_equal(rlt[0], (tensor_3d_1_np * 255).round())
np.testing.assert_almost_equal(rlt[1], (tensor_3d_2_np * 255).round())
rlt = tensor2img(tensor_3d_3, out_type=np.uint8, min_max=(0, 1))
tensor_3d_3_np = tensor_3d_3.squeeze().numpy()
np.testing.assert_almost_equal(rlt, (tensor_3d_3_np * 255).round())
# 2d
rlt = tensor2img(tensor_2d, out_type=np.uint8, min_max=(0, 1))
tensor_2d_np = tensor_2d.numpy()
np.testing.assert_almost_equal(rlt, (tensor_2d_np * 255).round())
rlt = tensor2img(tensor_2d, out_type=np.float32, min_max=(0, 1))
np.testing.assert_almost_equal(rlt, tensor_2d_np)
rlt = tensor2img(tensor_2d, out_type=np.float32, min_max=(0.1, 0.5))
tensor_2d_np = (np.clip(tensor_2d_np, 0.1, 0.5) - 0.1) / 0.4
np.testing.assert_almost_equal(rlt, tensor_2d_np)
def main():
autoN = auto_name()
parser = argparse.ArgumentParser()
parser.add_argument("--exper", default="./results/%s" % autoN, type=str)
parser.add_argument("--epoch", default=200, type=int)
parser.add_argument("--lr", default=0.005, type=float)
parser.add_argument("--bs", default=256, type=int)
parser.add_argument("--device", default="cuda:0", type=str)
parser.add_argument("--log_dir", default="./runs", type=str)
parser.add_argument("--num_workers", default=4, type=int)
parser.add_argument("--usefc", action="store_true")
parser.add_argument("--latent_num", default=100, type=int)
parser.add_argument("--kl_weight", default=0.001, type=float)
parser.add_argument("--loss_type", default="mse", type=str)
args = parser.parse_args()
device = torch.device(args.device)
kl_w = args.kl_weight
topilimage = ToPILImage()
# 读取数据集
if args.usefc:
transfer = Compose([ToTensor(), Lambda(lambda x: x.flatten())])
else:
transfer = Compose([ToTensor()])
train_data = MNIST("~/Datasets",
train=True,
download=True,
transform=transfer)
test_data = MNIST("~/Datasets",
train=False,
download=True,
transform=transfer)
train_dataloader = DataLoader(train_data,
batch_size=args.bs,
shuffle=True,
num_workers=args.num_workers)
test_dataloader = DataLoader(test_data,
batch_size=args.bs,
num_workers=args.num_workers)
# 构建模型
if args.usefc:
model = VanillaVAEfc(28 * 28, args.latent_num,
[1000, 500, 200]).to(device)
else:
model = VanillaVAEcnn(1, args.latent_num, [32, 64, 128, 256, 512])
optimizer = optim.Adam(model.parameters(), lr=args.lr)
# 训练
writer = SummaryWriter(osp.join(args.log_dir, osp.basename(args.exper)))
history = {
"train": {
"rec": [],
"kl": [],
"total": []
},
"test": {
"rec": [],
"kl": [],
"total": []
}
}
best = {"index": -1, "model": None, "loss": inf}
pilimgs = []
for e in tqdm(range(args.epoch), desc="Epoch: "):
writer.add_scalar("KL_weight", kl_w, e)
# train phase
loss_objs = [Loss() for _ in range(3)]
model.train()
for img, _ in train_dataloader:
img = img.to(device)
rec, mu, logvar = model(img)
loss, rec_loss, kl_loss = model.criterion(rec, img, mu, logvar,
kl_w, args.loss_type)
# 更新参数
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 记录训练的过程
bs = img.size(0)
for i, ls in enumerate([rec_loss, kl_loss, loss]):
loss_objs[i].add(ls, bs)
# epoch loss，并更新tensorboard
epoch_losses = [lo.value() for lo in loss_objs]
# minor_bar.set_description("Phase: test, Batch: ")
for i, s in enumerate(["rec", "kl", "total"]):
history["train"][s].append(epoch_losses[i])
writer.add_scalar("train/%s" % s, epoch_losses[i], e)
# test phase
loss_objs = [Loss() for _ in range(3)]
model.eval()
with torch.no_grad():
for img, _ in test_dataloader:
img = img.to(device)
rec, mu, logvar = model(img)
loss, rec_loss, kl_loss = model.criterion(
rec, img, mu, logvar, kl_w, args.loss_type)
# 记录训练的过程
bs = img.size(0)
for i, ls in enumerate([rec_loss, kl_loss, loss]):
loss_objs[i].add(ls, bs)
# epoch loss，并更新tensorboard
epoch_losses = [lo.value() for lo in loss_objs]
for i, s in enumerate(["rec", "kl", "total"]):
history["test"][s].append(epoch_losses[i])
writer.add_scalar("test/%s" % s, epoch_losses[i], e)
# 每5个epoch画一次图
if e % 5 == 0:
latents = torch.randn(64, args.latent_num, device=device)
gen_imgs = model.decode(latents)
if args.usefc:
gen_imgs = gen_imgs.reshape(64, 1, 28, 28)
writer.add_images("sampling", gen_imgs, e // 5)
pilimgs.append(topilimage(make_grid(gen_imgs).cpu()))
# best
if epoch_losses[-1] < best["loss"]:
best["index"] = e
best["loss"] = epoch_losses[-1]
best["model"] = deepcopy(model.state_dict())
# 保存结果
print("")
print("Best index: %d, Best Loss: %.4f" % (best["index"], best["loss"]))
if not osp.exists(args.exper):
mkdir(args.exper)
model.load_state_dict(best["model"])
torch.save(model, osp.join(args.exper, "model.pth"))
with open(osp.join(args.exper, "hist.json"), "w") as f:
json.dump(history, f)
best.pop("model")
with open(osp.join(args.exper, "best.json"), "w") as f:
json.dump(best, f)
with open(osp.join(args.exper, "args.json"), "w") as f:
json.dump(args.__dict__, f)
# 采样的图像创建gif
pilimgs[0].save("samples_%s.gif" % args.loss_type,
format="GIF",
append_images=pilimgs[1:],
save_all=True,
duration=500,
loop=0)
def get_batch_grid(batch,nrow):
return make_grid(batch).permute(1,2,0)
def train(opt):
# Init Model
generator = Generator().cuda()
discriminator = Discriminator().cuda()
discriminator.train()
# Load Dataset
train_dataset = MNISTDataset('train')
train_data_loader = MNISTDataloader('train', opt, train_dataset)
# test_dataset = MNISTDataset('test')
# test_data_loader = MNISTDataloader('test', opt, test_dataset)
# Set Optimizer
optim_gen = torch.optim.Adam(generator.parameters(), lr=0.0002)
optim_dis = torch.optim.Adam(discriminator.parameters(), lr=0.0002)
writer = SummaryWriter()
for epoch in range(opt.epoch):
for i in range(len(train_data_loader.data_loader)):
step = epoch * len(train_data_loader.data_loader) + i + 1
# load dataset only batch_size
image, label = train_data_loader.next_batch()
image = image.cuda()
# train discriminator
optim_dis.zero_grad()
noise = Variable(torch.randn(opt.batch_size, 100)).cuda()
gen = generator(noise)
validity_real = discriminator(image)
loss_dis_real = validity_real.mean()
(-loss_dis_real).backward()
validity_fake = discriminator(gen.detach())
loss_dis_fake = validity_fake.mean()
loss_dis_fake.backward()
# gradient penalty
eps = torch.rand(opt.batch_size, 1, 1, 1).cuda()
x_hat = eps * image.data + (1 - eps) * gen.data
x_hat.requires_grad = True
hat_predict = discriminator(x_hat)
grad_x_hat = grad(
outputs=hat_predict.sum(), inputs=x_hat, create_graph=True
)[0]
grad_penalty = ((grad_x_hat.view(grad_x_hat.size(0), -1).norm(2, dim=1) - 1) ** 2).mean()
grad_penalty = 10 * grad_penalty
grad_penalty.backward()
optim_dis.step()
loss_dis = -loss_dis_real + loss_dis_fake + grad_penalty
# train generator
generator.train()
optim_gen.zero_grad()
requires_grad(generator, True)
requires_grad(discriminator, False)
noise = Variable(torch.randn(opt.batch_size, 100)).cuda()
gen = generator(noise)
validity = discriminator(gen)
# print(validity)
loss_gen = validity.mean()
(-loss_gen).backward()
optim_gen.step()
requires_grad(generator, False)
requires_grad(discriminator, True)
writer.add_scalar('loss/gen/', loss_gen, step)
writer.add_scalar('loss/dis/', loss_dis, step)
if step % opt.display_step == 0:
writer.add_images('image', image[0][0], step, dataformats="HW")
writer.add_images('result', gen[0][0], step, dataformats="HW")
print('[Epoch {}] Total : {:.2} | G_loss : {:.2} | D_loss : {:.2}'.format(epoch + 1, loss_gen+loss_dis, loss_gen, loss_dis))
generator.eval()
z = Variable(torch.randn(9, 100)).cuda()
sample_images = generator(z)
grid = make_grid(sample_images, nrow=3, normalize=True)
writer.add_image('sample_image', grid, step)
torch.save(generator.state_dict(), 'checkpoint_{}.pt'.format(step))
def add_overlay(self, tag, embed, img, alpha=0.8, cmap='inferno', add_ref=None):
""" Adds to the summary the images of the input image (ref or search)
overlayed with the corresponding embedding or correlation map. It expect
tensors of with dimensions [C x H x W] or [B x C x H x W] if the tensor
has a batch dimension it takes the FIRST ELEMENT of the batch. The image
is displayed as fusion of the input image in grayscale and the overlay
in the chosen color_map, this fusion is controlled by the alpha factor.
In the case of the embeddings, since there are multiple feature
channels, we show each of them individually in a grid.
OBS: The colors represent relative values, where the peak color corresponds
to the maximum value in any given channel, so no direct value comparisons
can be made between epochs, only the relative distribution of neighboring
pixel values, (which should be enough, since we are mosly interested
in finding the maximum of a given correlation map)
Args:
tag: (str) The string identifying the image in tensorboard, images
with the same tag are grouped together with a slider, and are
indexed by epoch.
embed: (torch.Tensor) The tensor containing the embedding of an
input (ref or search image) or a correlation map (the final
output). The shape should be [B, C, H, W] or [B, H, W] for the
case of the correlation map.
img: (torch.Tensor) The image on top of which the embed is going
to be overlaid. Reference image embeddings should be overlaid
on top of reference images and search image embeddings as well
as the correlation maps should be overlaid on top of the search
images.
alpha: (float) A mixing variable, it controls how much of the final
embedding corresponds to the grayscale input image and how much
corresponds to the overlay. Alpha = 0, means there is no
overlay in the final image, only the input image. Conversely,
Alpha = 1 means there is only overlay. Adjust this value so
you can distinctly see the overlay details while still seeing
where it is in relation to the orignal image.
cmap: (str) The name of the colormap to be used with the overlay.
The colormaps are defined in the colormaps.py module, but values
include 'viridis' (greenish blue) and 'inferno' (yellowish red).
add_ref: (torch.Tensor) Optional. An additional reference image that
will be plotted to the side of the other images. Useful when
plotting correlation maps, because it lets the user see both
the search image and the reference that is used as the target.
``Example``
>>> summ_maker = SummaryMaker(os.path.join(exp_dir, 'tensorboard'), params,
model.upscale_factor)
...
>>> embed_ref = model.get_embedding(ref_img_batch)
>>> embed_srch = model.get_embedding(search_batch)
>>> output_batch = model.match_corr(embed_ref, embed_srch)
>>> batch_index = 0
>>> summ_maker.add_overlay("Ref_image_{}".format(tbx_index), embed_ref[batch_index], ref_img_batch[batch_index], cmap='inferno')
>>> summ_maker.add_overlay("Search_image_{}".format(tbx_index), embed_srch[batch_index], search_batch[batch_index], cmap='inferno')
>>> summ_maker.add_overlay("Correlation_map_{}".format(tbx_index), output_batch[batch_index], search_batch[batch_index], cmap='inferno')
"""
# TODO Add numbers in the final image to the feature channels.
# TODO Add the color bar showing the progression of values.
# If minibatch is given, take only the first image
# TODO let the user select the image? Loop on all images?
if len(embed.shape) == 4:
embed = embed[0]
if len(img.shape) == 4:
img = img[0]
# Normalize the image.
img = img - img.min()
img = img/img.max()
embed = cm.apply_cmap(embed, cmap=cmap)
# Get grayscale version of image by taking the weighted average of the channels
# as described in https://www.cs.virginia.edu/~vicente/recognition/notebooks/image_processing_lab.html#2.-Converting-to-Grayscale
R,G,B = img
img_gray = 0.21 * R + 0.72 * G + 0.07 * B
# Get the upscaled size of the embedding, so as to take into account
# the network's downscale caused by the stride.
upsc_size = (embed.shape[-1] - 1) * self.up_factor + 1
embed = F.interpolate(embed, upsc_size, mode='bilinear',
align_corners=False)
# Pad the embedding with zeros to match the image dimensions. We pad
# all 4 corners equally to keep the embedding centered.
tot_pad = img.shape[-1] - upsc_size
# Sanity check 1. The amount of padding must be equal on all sides, so
# the total padding on any dimension must be an even integer.
assert tot_pad % 2 == 0, "The embed or image dimensions are incorrect."
pad = int(tot_pad/2)
embed = F.pad(embed, (pad, pad, pad, pad), 'constant', 0)
# Sanity check 2, the size of the embedding in the (H, w) dimensions
# matches the size of the image.
assert embed.shape[-2:] == img.shape[-2:], ("The embedding overlay "
"and image dimensions "
"do not agree.")
final_imgs = alpha * embed + (1-alpha) * img_gray
# The embedding_channel (or feature channel) dimension is treated like
# a batch dimension, so the grid shows each individual embeding
# overlayed with the input image. Plus the original image is also shown.
# If add_ref is used the ref image is the first to be shown.
img = img.unsqueeze(0)
final_imgs = torch.cat((img, final_imgs))
if add_ref is not None:
# Pads the image if necessary
pad = int((img.shape[-1] - add_ref.shape[-1])//2)
add_ref = F.pad(add_ref, (pad, pad, pad, pad), 'constant', 0)
add_ref = add_ref.unsqueeze(0)
final_imgs = torch.cat((add_ref, final_imgs))
final_imgs = make_grid(final_imgs, nrow=6)
self.writer_val.add_image(tag, final_imgs, self.epoch)
actor_model,
transition_model,
encoder,
belief,
posterior_state,
action,
observation.to(device=args.device),
deterministic=True)
# total_rewards += reward.numpy()
total_rewards += reward
if not args.symbolic: # Collect real vs. predicted frames for video
video_frames.append(
make_grid(torch.cat([
observation,
observation_model(belief, posterior_state).cpu()
],
dim=3) + 0.5,
nrow=5).numpy()) # Decentre
observation = next_observation
# if done.sum().item() == args.test_episodes:
# pbar.close()
# break
if done:
pbar.close()
break
# Update and plot reward metrics (and write video if applicable) and save metrics
metrics['test_episodes'].append(episode)
# metrics['test_rewards'].append(total_rewards.tolist())
metrics['test_rewards'].append(total_rewards)
lineplot(metrics['test_episodes'], metrics['test_rewards'],
im = portrait_transform(im)
im = im.to(device)
im = im.unsqueeze(0)
if opt.model_type == 'can32':
model = CAN(n_channels=32)
if opt.model_type == 'sandocan32':
model = SandOCAN()
if opt.model_type == 'unet':
model = UNet()
assert model
models_path = [
f for f in os.listdir(opt.final_dir) if f.startswith(opt.model_type)
]
#print(models_path)
images = im
for model_name in models_path:
print('Loading model' + model_name)
model.load_state_dict(
torch.load(os.path.join(opt.final_dir, model_name),
map_location=lambda storage, loc: storage))
model.to(device)
images = torch.cat((images, model(im)))
names = [' '.join(m.split('_')[1:-1]) for m in models_path]
filename = opt.model_type + 'actual_' + '_'.join(names) + '.png'
save_image(images, filename, normalize=True, range=(-1, 1))
grid = make_grid(images, nrow=1, normalize=True, range=(-1, 1))
writer.add_image(filename, grid)
def validate(val_loader, net, criterion, optimizer, epoch, train_args, visualize):
# the following code is written assuming that batch size is 1
net.eval()
val_loss = AverageMeter()
gts_all = np.zeros((len(val_loader), args['shorter_size'], 2 * args['shorter_size']), dtype=int)
predictions_all = np.zeros((len(val_loader), args['shorter_size'], 2 * args['shorter_size']), dtype=int)
for vi, data in enumerate(val_loader):
input, gt, slices_info = data
assert len(input.size()) == 5 and len(gt.size()) == 4 and len(slices_info.size()) == 3
input.transpose_(0, 1)
gt.transpose_(0, 1)
slices_info.squeeze_(0)
assert input.size()[3:] == gt.size()[2:]
count = torch.zeros(args['shorter_size'], 2 * args['shorter_size']).cuda()
output = torch.zeros(voc.num_classes, args['shorter_size'], 2 * args['shorter_size']).cuda()
slice_batch_pixel_size = input.size(1) * input.size(3) * input.size(4)
for input_slice, gt_slice, info in zip(input, gt, slices_info):
input_slice = Variable(input_slice).cuda()
gt_slice = Variable(gt_slice).cuda()
output_slice = net(input_slice)
assert output_slice.size()[2:] == gt_slice.size()[1:]
assert output_slice.size()[1] == voc.num_classes
output[:, info[0]: info[1], info[2]: info[3]] += output_slice[0, :, :info[4], :info[5]].data
gts_all[vi, info[0]: info[1], info[2]: info[3]] += gt_slice[0, :info[4], :info[5]].data.cpu().numpy()
count[info[0]: info[1], info[2]: info[3]] += 1
val_loss.update(criterion(output_slice, gt_slice).data[0], slice_batch_pixel_size)
output /= count
gts_all[vi, :, :] /= count.cpu().numpy().astype(int)
predictions_all[vi, :, :] = output.max(0)[1].squeeze_(0).cpu().numpy()
print('validating: %d / %d' % (vi + 1, len(val_loader)))
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, voc.num_classes)
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr'])
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(gts_all, predictions_all)):
gt_pil = voc.colorize_mask(data[0])
predictions_pil = voc.colorize_mask(data[1])
if train_args['val_save_to_img_file']:
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=2, padding=5)
writer.add_image(snapshot_name, val_visual)
print('-----------------------------------------------------------------------------------------------------------')
print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]' % (
train_args['best_record']['val_loss'], train_args['best_record']['acc'], train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'], train_args['best_record']['fwavacc'], train_args['best_record']['epoch']))
print('-----------------------------------------------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
net.train()
return val_loss.avg
print(z_interp.shape)
inputs_row = torch.zeros((n_interp, ) + x_interp.shape[1:])
inputs_row[0] = inputs[idx_start]
inputs_row[-1] = inputs[idx_end]
x_interp = torch.cat([x_interp.cpu(), inputs_row], dim=0)
# plt.imshow(make_grid(inputs.cpu(), normalize=True).permute(1, 2, 0))
# plt.show()
# plt.imshow(make_grid(x_recon.cpu(), normalize=True).permute(1, 2, 0))
# plt.show()
plt.imshow(
make_grid(x_interp.cpu(), normalize=True, nrow=n_interp).permute(1, 2, 0))
plt.show()
# for epoch in range(args.n_epochs):
# for i, (inputs, _) in enumerate(dataloader):
# inputs = inputs.to(device)
# optimizer.zero_grad()
# mu, logvar = netE.encode(inputs)
# z = netE.reparameterize(mu, logvar)
# x_recon = netD(z)
# loss = criterion(mu, logvar, x_recon, inputs)
# loss.backward()
# optimizer.step()
def showLandmarksBatch(sampleBatched):
imagesBatch, landmarksBatch = sampleBatched['image'], sampleBatched['landmarks']
batchSize = imagesBatch.shape[0]
imSize = imagesBatch.size(2)
grid = utils.make_grid(imagesBatch)
plt.imshow(grid.numpy().transpose((1, 2, 0)))
def validate(val_loader, net, criterion, optimizer, epoch, train_args, restore, visualize):
net.eval()
val_loss = AverageMeter()
inputs_all, gts_all, predictions_all = [], [], []
for vi, data in enumerate(val_loader):
inputs, gts = data
N = inputs.size(0)
inputs = Variable(inputs, volatile=True).cuda()
gts = Variable(gts, volatile=True).cuda()
outputs = net(inputs)
predictions = outputs.data.max(1)[1].squeeze_(1).squeeze_(0).cpu().numpy()
val_loss.update(criterion(outputs, gts).data[0] / N, N)
if random.random() > train_args['val_img_sample_rate']:
inputs_all.append(None)
else:
inputs_all.append(inputs.data.squeeze_(0).cpu())
gts_all.append(gts.data.squeeze_(0).cpu().numpy())
predictions_all.append(predictions)
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, voc.num_classes)
if mean_iu > train_args['best_record']['mean_iu']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr']
)
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0])
gt_pil = voc.colorize_mask(data[1])
predictions_pil = voc.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(input_pil.convert('RGB')), visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
print('--------------------------------------------------------------------')
print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]' % (
train_args['best_record']['val_loss'], train_args['best_record']['acc'], train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'], train_args['best_record']['fwavacc'], train_args['best_record']['epoch']))
print('--------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
writer.add_scalar('lr', optimizer.param_groups[1]['lr'], epoch)
net.train()
return val_loss.avg
def main():
# Load dataset
print('Loading dataset ...\n')
dataset_train = Dataset(train=True)
dataset_val = Dataset(train=False)
loader_train = DataLoader(dataset=dataset_train,
num_workers=4,
batch_size=opt.batchSize,
shuffle=True)
print("# of training samples: %d\n" % int(len(dataset_train)))
# Build model
net = DnCNN(channels=1, num_of_layers=opt.num_of_layers)
net.apply(weights_init_kaiming)
criterion = nn.MSELoss(size_average=False)
# Move to GPU
device_ids = [0]
model = nn.DataParallel(net, device_ids=device_ids).cuda()
criterion.cuda()
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=opt.lr)
# training
writer = SummaryWriter(opt.outf)
step = 0
noiseL_B = [0, 55] # ingnored when opt.mode=='S'
for epoch in range(opt.epochs):
if epoch < opt.milestone:
current_lr = opt.lr
else:
current_lr = opt.lr / 10.
# set learning rate
for param_group in optimizer.param_groups:
param_group["lr"] = current_lr
print('learning rate %f' % current_lr)
# train
for i, data in enumerate(loader_train, 0):
# training step
model.train()
model.zero_grad()
optimizer.zero_grad()
img_train = data
if opt.mode == 'S':
noise = torch.FloatTensor(img_train.size()).normal_(
mean=0, std=opt.noiseL / 255.)
if opt.mode == 'B':
noise = torch.zeros(img_train.size())
stdN = np.random.uniform(noiseL_B[0],
noiseL_B[1],
size=noise.size()[0])
for n in range(noise.size()[0]):
sizeN = noise[0, :, :, :].size()
noise[n, :, :, :] = torch.FloatTensor(sizeN).normal_(
mean=0, std=stdN[n] / 255.)
imgn_train = img_train + noise
img_train, imgn_train = Variable(img_train.cuda()), Variable(
imgn_train.cuda())
noise = Variable(noise.cuda())
out_train = model(imgn_train)
loss = criterion(out_train, noise) / (imgn_train.size()[0] * 2)
loss.backward()
optimizer.step()
# results
model.eval()
out_train = torch.clamp(imgn_train - model(imgn_train), 0., 1.)
psnr_train = batch_PSNR(out_train, img_train, 1.)
print(
"[epoch %d][%d/%d] loss: %.4f PSNR_train: %.4f" %
(epoch + 1, i + 1, len(loader_train), loss.item(), psnr_train))
# if you are using older version of PyTorch, you may need to change loss.item() to loss.data[0]
if step % 10 == 0:
# Log the scalar values
writer.add_scalar('loss', loss.item(), step)
writer.add_scalar('PSNR on training data', psnr_train, step)
step += 1
## the end of each epoch
model.eval()
# validate
psnr_val = 0
for k in range(len(dataset_val)):
img_val = torch.unsqueeze(dataset_val[k], 0)
noise = torch.FloatTensor(img_val.size()).normal_(
mean=0, std=opt.val_noiseL / 255.)
imgn_val = img_val + noise
img_val, imgn_val = Variable(img_val.cuda()), Variable(
imgn_val.cuda())
out_val = torch.clamp(imgn_val - model(imgn_val), 0., 1.)
psnr_val += batch_PSNR(out_val, img_val, 1.)
psnr_val /= len(dataset_val)
print("\n[epoch %d] PSNR_val: %.4f" % (epoch + 1, psnr_val))
writer.add_scalar('PSNR on validation data', psnr_val, epoch)
# log the images
out_train = torch.clamp(imgn_train - model(imgn_train), 0., 1.)
Img = utils.make_grid(img_train.data,
nrow=8,
normalize=True,
scale_each=True)
Imgn = utils.make_grid(imgn_train.data,
nrow=8,
normalize=True,
scale_each=True)
Irecon = utils.make_grid(out_train.data,
nrow=8,
normalize=True,
scale_each=True)
writer.add_image('clean image', Img, epoch)
writer.add_image('noisy image', Imgn, epoch)
writer.add_image('reconstructed image', Irecon, epoch)
# save model
torch.save(model.state_dict(), os.path.join(opt.outf, 'net.pth'))
def train(self):
phase = 'train'
since = time.time()
best_loss = float('inf')
tensorboardX_iter_count = 0
for epoch in range(self.total_epoch_num):
print('\nEpoch {}/{}'.format(epoch + 1, self.total_epoch_num))
print('-' * 10)
fn = open(self.train_log, 'a')
fn.write('\nEpoch {}/{}\n'.format(epoch + 1, self.total_epoch_num))
fn.write('--' * 5 + '\n')
fn.close()
self._set_models_train(['depthEstModel'])
iterCount = 0
for sample_dict in self.dataloaders_xLabels_joint:
imageListReal, depthListReal = sample_dict['real']
imageListSyn, depthListSyn = sample_dict['syn']
imageListSyn = imageListSyn.to(self.device)
depthListSyn = depthListSyn.to(self.device)
imageListReal = imageListReal.to(self.device)
depthListReal = depthListReal.to(self.device)
valid_mask = (depthListReal > -1.) # remove undefined regions
with torch.set_grad_enabled(phase == 'train'):
total_loss = 0.
self.depth_optimizer.zero_grad()
real_depth_loss = self.compute_depth_loss(
imageListReal, depthListReal, self.depthEstModel,
valid_mask)
syn_depth_loss = self.compute_depth_loss(
imageListSyn, depthListSyn, self.depthEstModel)
total_loss += (real_depth_loss + syn_depth_loss)
if self.use_apex:
with amp.scale_loss(
total_loss,
self.depth_optimizer) as total_loss_scaled:
total_loss_scaled.backward()
else:
total_loss.backward()
self.depth_optimizer.step()
iterCount += 1
if self.use_tensorboardX:
self.train_display_freq = len(
self.dataloaders_xLabels_joint)
nrow = imageListReal.size()[0]
if tensorboardX_iter_count % self.train_display_freq == 0:
pred_depth_real = self.depthEstModel(imageListReal)[-1]
tensorboardX_grid_real_rgb = make_grid(imageListReal,
nrow=nrow,
normalize=True,
range=(-1.0,
1.0))
self.train_SummaryWriter.add_image(
'real rgb images', tensorboardX_grid_real_rgb,
tensorboardX_iter_count)
tensorboardX_depth_concat = torch.cat(
(depthListReal, pred_depth_real), dim=0)
tensorboardX_grid_real_depth = make_grid(
tensorboardX_depth_concat,
nrow=nrow,
normalize=True,
range=(-1.0, 1.0))
self.train_SummaryWriter.add_image(
'real depth and depth prediction',
tensorboardX_grid_real_depth,
tensorboardX_iter_count)
# add loss values
loss_val_list = [
total_loss, real_depth_loss, syn_depth_loss
]
loss_name_list = [
'total_loss', 'real_depth_loss', 'syn_depth_loss'
]
self.write_2_tensorboardX(self.train_SummaryWriter,
loss_val_list,
name=loss_name_list,
mode='scalar',
count=tensorboardX_iter_count)
tensorboardX_iter_count += 1
if iterCount % 20 == 0:
loss_summary = '\t{}/{} total_loss: {:.7f}, real_depth_loss: {:.7f}, syn_depth_loss: {:.7f}'.format(
iterCount, len(self.dataloaders_xLabels_joint),
total_loss, real_depth_loss, syn_depth_loss)
print(loss_summary)
fn = open(self.train_log, 'a')
fn.write(loss_summary)
fn.close()
# take step in optimizer
for scheduler in self.scheduler_list:
scheduler.step()
for optim in self.optim_name:
lr = getattr(self, optim).param_groups[0]['lr']
lr_update = 'Epoch {}/{} finished: {} learning rate = {:.7f}'.format(
epoch + 1, self.total_epoch_num, optim, lr)
print(lr_update)
fn = open(self.train_log, 'a')
fn.write(lr_update)
fn.close()
if (epoch + 1) % self.save_steps == 0:
self.save_models(self.model_name, mode=epoch + 1)
self.evaluate(epoch + 1)
time_elapsed = time.time() - since
print('\nTraining complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
fn = open(self.train_log, 'a')
fn.write('\nTraining complete in {:.0f}m {:.0f}s\n'.format(
time_elapsed // 60, time_elapsed % 60))
fn.close()
best_model_summary = '\nOverall best model is epoch {}'.format(
self.EVAL_best_model_epoch)
print(best_model_summary)
print(self.EVAL_all_results[str(self.EVAL_best_model_epoch)])
fn = open(self.evaluate_log, 'a')
fn.write(best_model_summary + '\n')
fn.write(self.EVAL_all_results[str(self.EVAL_best_model_epoch)])
fn.close()
# Output training stats
if i % 10 == 0:
print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
% (epoch, num_epochs, i, len(dataloader),
errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
# Save Losses for plotting later
G_losses.append(errG.item())
D_losses.append(errD.item())
# Check how the generator is doing by saving G's output on fixed_noise
if (iters % 500 == 0) or ((epoch == num_epochs-1) and (i == len(dataloader)-1)):
with torch.no_grad():
fake = netG(fixed_noise).detach().cpu()
img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
iters += 1
plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses,label="G")
plt.plot(D_losses,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()
#%%capture
fig = plt.figure(figsize=(4, 4))
def train_summaries(self, batch, epoch):
"""Tensorboard logging."""
if self.options.stage == 'dp':
dtype = self.vis_data['pred_dp'].dtype
rend_imgs = []
vis_size = self.vis_data['pred_dp'].shape[0]
# Do visualization for the first 4 images of the batch
for i in range(vis_size):
img = self.vis_data['image'][i].cpu().numpy().transpose(
1, 2, 0)
H, W, C = img.shape
rend_img = img.transpose(2, 0, 1)
gt_dp = self.vis_data['gt_dp'][i]
gt_dp = torch.nn.functional.interpolate(gt_dp[None, :],
size=[H, W])[0]
# gt_dp = torch.cat((gt_dp, gt_dp.new_ones(1, H, W)), dim=0).cpu().numpy()
gt_dp = gt_dp.cpu().numpy()
rend_img = np.concatenate((rend_img, gt_dp), axis=2)
pred_dp = self.vis_data['pred_dp'][i]
pred_dp[0] = (pred_dp[0] > 0.5).type(dtype)
pred_dp[1:] = pred_dp[1:] * pred_dp[0]
pred_dp = torch.nn.functional.interpolate(pred_dp[None, :],
size=[H, W])[0]
pred_dp = pred_dp.cpu().numpy()
rend_img = np.concatenate((rend_img, pred_dp), axis=2)
# import matplotlib.pyplot as plt
# plt.imshow(rend_img.transpose([1, 2, 0]))
rend_imgs.append(torch.from_numpy(rend_img))
rend_imgs = make_grid(rend_imgs, nrow=1)
self.summary_writer.add_image('imgs', rend_imgs, self.step_count)
else:
gt_keypoints_2d = self.vis_data['gt_joint'].cpu().numpy()
pred_vertices = self.vis_data['pred_vert']
pred_keypoints_2d = self.vis_data['pred_joint']
pred_camera = self.vis_data['pred_cam']
dtype = pred_camera.dtype
rend_imgs = []
vis_size = pred_vertices.shape[0]
# Do visualization for the first 4 images of the batch
for i in range(vis_size):
img = self.vis_data['image'][i].cpu().numpy().transpose(
1, 2, 0)
H, W, C = img.shape
# Get LSP keypoints from the full list of keypoints
gt_keypoints_2d_ = gt_keypoints_2d[i, self.to_lsp]
pred_keypoints_2d_ = pred_keypoints_2d.cpu().numpy()[
i, self.to_lsp]
vertices = pred_vertices[i].cpu().numpy()
cam = pred_camera[i].cpu().numpy()
# Visualize reconstruction and detected pose
rend_img = visualize_reconstruction(img, self.options.img_res,
gt_keypoints_2d_, vertices,
pred_keypoints_2d_, cam,
self.renderer)
rend_img = rend_img.transpose(2, 0, 1)
if 'gt_vert' in self.vis_data.keys():
rend_img2 = vis_mesh(
img,
self.vis_data['gt_vert'][i].cpu().numpy(),
cam,
self.renderer,
color='blue')
rend_img2 = rend_img2.transpose(2, 0, 1)
rend_img = np.concatenate((rend_img, rend_img2), axis=2)
gt_dp = self.vis_data['gt_dp'][i]
gt_dp = torch.nn.functional.interpolate(gt_dp[None, :],
size=[H, W])[0]
gt_dp = gt_dp.cpu().numpy()
# gt_dp = torch.cat((gt_dp, gt_dp.new_ones(1, H, W)), dim=0).cpu().numpy()
rend_img = np.concatenate((rend_img, gt_dp), axis=2)
pred_dp = self.vis_data['pred_dp'][i]
pred_dp[0] = (pred_dp[0] > 0.5).type(dtype)
pred_dp[1:] = pred_dp[1:] * pred_dp[0]
pred_dp = torch.nn.functional.interpolate(pred_dp[None, :],
size=[H, W])[0]
pred_dp = pred_dp.cpu().numpy()
rend_img = np.concatenate((rend_img, pred_dp), axis=2)
# import matplotlib.pyplot as plt
# plt.imshow(rend_img.transpose([1, 2, 0]))
rend_imgs.append(torch.from_numpy(rend_img))
rend_imgs = make_grid(rend_imgs, nrow=1)
uv_maps = []
for i in range(vis_size):
uv_temp = torch.cat(
(self.vis_data['pred_uv'][i], self.vis_data['gt_uv'][i]),
dim=1)
uv_maps.append(uv_temp.permute(2, 0, 1))
uv_maps = make_grid(uv_maps, nrow=1)
uv_maps = uv_maps.abs()
uv_maps = uv_maps / uv_maps.max()
# Save results in Tensorboard
self.summary_writer.add_image('imgs', rend_imgs, self.step_count)
self.summary_writer.add_image('uv_maps', uv_maps, self.step_count)
for key in self.loss_item.keys():
self.summary_writer.add_scalar('loss_' + key, self.loss_item[key],
self.step_count)
def tensor2image(x):
tensor = x2num.make_np(vutils.make_grid(x.data[:64], normalize=True))
xtensors = xutils.convert_to_HWC(tensor, 'CHW')
plt.imshow(xtensors)
plt.show()
if epoch == 0:
best_val = valid_loss
elif (valid_loss < best_val):
save_model(epoch, G)
pay = 0
best_val = valid_loss
pay = pay + 1
if (pay == 100):
break
print(valid_loss)
train_loss_list.append(train_loss)
valid_loss_list.append(valid_loss)
rec_imgs, z = G(fixed_batch)
show_and_save('results/Input_epoch_%d.png' % epoch,
make_grid((fixed_batch.data[:, 2:3, :, :]).cpu(), 8))
show_and_save('results/rec_epoch_%d.png' % epoch,
make_grid((rec_imgs.data[:, 2:3, :, :]).cpu(), 8))
show_and_save(
'results/Error_epoch_%d.png' % epoch,
make_grid((fixed_batch.data[:, 2:3, :, :] -
rec_imgs.data[:, 2:3, :, :]).cpu(), 8))
#localtime = time.asctime( time.localtime(time.time()) )
#D_real_list_np=(D_real_list).to('cpu')
save_model(epoch, G)
plt.plot(train_loss_list, label="train loss")
plt.plot(valid_loss_list, label="validation loss")
plt.legend()
plt.show()
flag = 0
# flag = 1
if flag:
z_idx = 0
single_noise = torch.randn(1, nz, 1, 1, device=device)
for i in range(num_img):
add_noise = single_noise
add_noise = add_noise[0, z_idx, 0, 0] + i * 0.01
fixed_noise[i, ...] = add_noise
# step 2: model
net_g = Generator(nz=nz, ngf=ngf, nc=nc)
# net_d = Discriminator(nc=nc, ndf=ndf)
checkpoint = torch.load(path_checkpoint, map_location="cpu")
state_dict_g = checkpoint["g_model_state_dict"]
state_dict_g = remove_module(state_dict_g)
net_g.load_state_dict(state_dict_g)
net_g.to(device)
# net_d.load_state_dict(checkpoint["d_model_state_dict"])
# net_d.to(device)
# step3: inference
with torch.no_grad():
fake_data = net_g(fixed_noise).detach().cpu()
img_grid = vutils.make_grid(fake_data, padding=2, normalize=True).numpy()
img_grid = np.transpose(img_grid, (1, 2, 0))
plt.imshow(img_grid)
plt.show()
for i, imgs in enumerate(dataloader):
# Configure model input
lr_imgs = imgs[0].to(device)
hr_imgs = imgs[1].to(device)
### train Generator
pre_G_optimizer.zero_grad()
G_loss = criterion_content(G(lr_imgs), hr_imgs)
G_loss.backward()
G_optimizer.step()
if epoch % 20 == 0:
sys.stdout.write(
"[Epoch %d/%d] [Batch %d/%d] [G loss: %f]" %
(epoch, pre_epochs, i, len(dataloader), G_loss.item()))
fake_imgs_ = make_grid(G(lr_imgs), nrow=1, normalize=True)
save_image(fake_imgs_,
"images/train/pre%d.png" % (epoch + 200),
normalize=False)
fake_imgs_ = make_grid(lr_imgs, nrow=1, normalize=True)
save_image(fake_imgs_,
"images/train/pre%d_lr.png" % (epoch + 200),
normalize=False)
G.eval()
# validation
with torch.no_grad():
valing_results = {'mse': 0, 'psnr': 0, 'batch_sizes': 0}
val_images = []
val_dataset = ValDataset(4)
val_dataloader = DataLoader(val_dataset,
# - transforms：提供常用的数据预处理操作，主要包括对Tensor以及PIL Image对象的操作。
#我们下面来看一下如何才能加载已经训练好的模型
from torchvision import models
#resnet34=models.squeezenet1_1(pretrained=True,num_classes=1000)
#resnet34.fc=nn.Linear(512,10)#然后修改其全连接层，使其变为10分类的问题
#下载的未知为/home/dongxinge/.torch/models/squeezenet1_1-f364aa15.pth
from torchvision import datasets
#dataset=datasets.MNIST('data/',download=True,train=False,transform=transforms)
#data/为数据集下载的位置
#train=False为下载的数据为测试集
#transform为下载的数据进行转换操作
from torchvision import transforms
to_pil = transforms.ToPILImage() #构建一个transform转换操作
a = to_pil(torch.randn(3, 64, 64)) #对torch.randn(3,64,64)进行转换操作。此时它还是图片，并不是张量
#torchvison的两个常用的函数，
#make_grid,它能够将多张图片拼接成一个网格中
#save_img它能将tensor保存成图片
dataloader = DataLoader(dataset, shuffle=True, batch_size=16)
from torchvision.utils import make_grid, save_image
dataiter = iter(dataloader)
#print(len(next(dataiter)[0]))8,因为一共有8张图片，所以虽然batch_size为16，但是只能取出8张来
img = make_grid(next(dataiter)[0], 4) #拼接成4*4的网格图片，并且会将图片转成3通道的
#print(img)<PIL.Image.Image image mode=RGB size=906x454 at 0x7F102350C240>
a = to_img(img) #将tensor转成图片
#print(a)#<PIL.Image.Image image mode=RGB size=906x454 at 0x7F2C637872B0>
#save_image(img, 'a.png')#保存图片，名称是a.png
#Image.open('a.png')
def validate(val_loader, net, criterion, optimizer, epoch, train_args, restore, visualize):
net.eval()
val_loss = AverageMeter()
inputs_all, gts_all, predictions_all = [], [], []
for vi, data in enumerate(val_loader):
inputs, gts = data
N = inputs.size(0)
inputs = Variable(inputs, volatile=True).cuda()
gts = Variable(gts, volatile=True).cuda()
outputs = net(inputs)
predictions = outputs.data.max(1)[1].squeeze_(1).squeeze_(0).cpu().numpy()
val_loss.update(criterion(outputs, gts).data[0] / N, N)
if random.random() > train_args['val_img_sample_rate']:
inputs_all.append(None)
else:
inputs_all.append(inputs.data.squeeze_(0).cpu())
gts_all.append(gts.data.squeeze_(0).cpu().numpy())
predictions_all.append(predictions)
acc, acc_cls, mean_iu, fwavacc = evaluate(predictions_all, gts_all, voc.num_classes)
if mean_iu > train_args['best_record']['mean_iu']:
train_args['best_record']['val_loss'] = val_loss.avg
train_args['best_record']['epoch'] = epoch
train_args['best_record']['acc'] = acc
train_args['best_record']['acc_cls'] = acc_cls
train_args['best_record']['mean_iu'] = mean_iu
train_args['best_record']['fwavacc'] = fwavacc
snapshot_name = 'epoch_%d_loss_%.5f_acc_%.5f_acc-cls_%.5f_mean-iu_%.5f_fwavacc_%.5f_lr_%.10f' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc, optimizer.param_groups[1]['lr']
)
torch.save(net.state_dict(), os.path.join(ckpt_path, exp_name, snapshot_name + '.pth'))
torch.save(optimizer.state_dict(), os.path.join(ckpt_path, exp_name, 'opt_' + snapshot_name + '.pth'))
if train_args['val_save_to_img_file']:
to_save_dir = os.path.join(ckpt_path, exp_name, str(epoch))
check_mkdir(to_save_dir)
val_visual = []
for idx, data in enumerate(zip(inputs_all, gts_all, predictions_all)):
if data[0] is None:
continue
input_pil = restore(data[0])
gt_pil = voc.colorize_mask(data[1])
predictions_pil = voc.colorize_mask(data[2])
if train_args['val_save_to_img_file']:
input_pil.save(os.path.join(to_save_dir, '%d_input.png' % idx))
predictions_pil.save(os.path.join(to_save_dir, '%d_prediction.png' % idx))
gt_pil.save(os.path.join(to_save_dir, '%d_gt.png' % idx))
val_visual.extend([visualize(input_pil.convert('RGB')), visualize(gt_pil.convert('RGB')),
visualize(predictions_pil.convert('RGB'))])
val_visual = torch.stack(val_visual, 0)
val_visual = vutils.make_grid(val_visual, nrow=3, padding=5)
writer.add_image(snapshot_name, val_visual)
print('--------------------------------------------------------------------')
print('[epoch %d], [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f]' % (
epoch, val_loss.avg, acc, acc_cls, mean_iu, fwavacc))
print('best record: [val loss %.5f], [acc %.5f], [acc_cls %.5f], [mean_iu %.5f], [fwavacc %.5f], [epoch %d]' % (
train_args['best_record']['val_loss'], train_args['best_record']['acc'], train_args['best_record']['acc_cls'],
train_args['best_record']['mean_iu'], train_args['best_record']['fwavacc'], train_args['best_record']['epoch']))
print('--------------------------------------------------------------------')
writer.add_scalar('val_loss', val_loss.avg, epoch)
writer.add_scalar('acc', acc, epoch)
writer.add_scalar('acc_cls', acc_cls, epoch)
writer.add_scalar('mean_iu', mean_iu, epoch)
writer.add_scalar('fwavacc', fwavacc, epoch)
writer.add_scalar('lr', optimizer.param_groups[1]['lr'], epoch)
net.train()
return val_loss.avg
def show_per_batch_transform(self, batch: Any, _) -> None:
images = batch[0]["input"]
image = make_grid(images, nrow=2)
image = T.to_pil_image(image, 'RGB')
image.show()
dis_output_fake_v = net_discr(gen_output_v.detach())
dis_loss = objective(dis_output_true_v, true_labels_v) + \
objective(dis_output_fake_v, fake_labels_v)
dis_loss.backward()
dis_optimizer.step()
dis_losses.append(dis_loss.item())
# train generator
gen_optimizer.zero_grad()
dis_output_v = net_discr(gen_output_v)
gen_loss_v = objective(dis_output_v, true_labels_v)
gen_loss_v.backward()
gen_optimizer.step()
gen_losses.append(gen_loss_v.item())
iter_no += 1
if iter_no % REPORT_EVERY_ITER == 0:
log.info("Iter %d: gen_loss=%.3e, dis_loss=%.3e", iter_no,
np.mean(gen_losses), np.mean(dis_losses))
writer.add_scalar("gen_loss", np.mean(gen_losses), iter_no)
writer.add_scalar("dis_loss", np.mean(dis_losses), iter_no)
gen_losses = []
dis_losses = []
if iter_no % SAVE_IMAGE_EVERY_ITER == 0:
writer.add_image(
"fake", vutils.make_grid(gen_output_v.data[:64],
normalize=True), iter_no)
writer.add_image(
"real", vutils.make_grid(batch_v.data[:64], normalize=True),
iter_no)
def _write_images(self, name:str, images:[Tensor])->None:
"Writes list of images as tensors to Tensorboard."
tag = self.ds_type.name + ' ' + name
self.tbwriter.add_image(tag=tag, img_tensor=vutils.make_grid(images, normalize=True), global_step=self.iteration)
