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 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)