def train(args): # tensorboard run_name = "./runs/run-classifier_batch_" + str(args.batch_size) \ + "_epochs_" + str(args.epochs) + "_" + args.log_message configure(run_name) net = Net() mnistmTrainSet = ds.mnistmTrainingDataset(text_file=args.dataset_list) mnistmTrainLoader = torch.utils.data.DataLoader(mnistmTrainSet, batch_size=args.batch_size, shuffle=True, num_workers=2) loss_fn = nn.CrossEntropyLoss() optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9) # put on gpu if available device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net.to(device) epoch = [0] # load prev model tio.load_model(model=net, optimizer=optimizer, epoch=epoch) epoch = epoch[0] while epoch < args.epochs: for i, sample_batched in enumerate(mnistmTrainLoader, 0): input_batch = f.pad(sample_batched['image'].float(), (2, 2, 2, 2)) input_batch = input_batch.to(device) optimizer.zero_grad() output = net(input_batch) loss = loss_fn(output, sample_batched['labels'].to(device)) loss.backward() optimizer.step() # print statistics if i % 50 == 0: count = int( epoch * math.floor(len(mnistmTrainSet) / (args.batch_size * 200)) + (i / 200)) print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, loss.item())) log_value('loss', loss.item(), count) _, ind = output.max(1) name = 'pred_' + str(ind[0]) sample_image = sample_batched['image'][0] log_images(name, sample_image, count) # save model tio.save_model(epoch=epoch, model=net, optimizer=optimizer) epoch = epoch + 1
def evaluate(self, test_loader, metrics, epoch=0): N = 0 metric_results = [0] * len(metrics) classified = [] misclassified = [] self.model.eval() with torch.no_grad(): for step, (data, targets) in enumerate(test_loader): # import matplotlib.pyplot as plt # img = plt.imshow(data.numpy()[0,1,:]) # plt.show() try: data = data.to(self.device) targets = targets.to(self.device) except: data = data.to("cpu") targets = targets.to("cpu") batch_size = data.size(0) outputs = self.model(data) if self.images_plot_count > 0: _, pred = outputs.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(targets.view(1, -1).expand_as(pred)).cpu().numpy()[0] data = data.cpu().numpy() classified += list(data[correct.astype(bool)]) misclassified += list(data[(1-correct).astype(bool)]) if len(classified) > self.images_plot_count: classified = random.sample(classified, self.images_plot_count) if len(misclassified) > self.images_plot_count: misclassified = random.sample(misclassified, self.images_plot_count) # print('Valid:', ' '.join(str(outputs).split('\n')[0:2])) # print('Shape:', outputs.shape, 'Sums', str(outputs.cpu().numpy().sum(1)).replace('\n', '')) for i, metric in enumerate(metrics): metric_results[i] += metric(outputs.data, targets.data) * batch_size N += batch_size if self.images_plot_count > 0: import tensorboard_logger as tl tl.log_images('Valid_Classified/Image', classified, step=epoch) tl.log_images('Valid_Misclassified/Image', misclassified, step=epoch) self.model.train() return [res / N for res in metric_results]
def _log_reconstruction_imgs(title, sess, dataset, inp, out, epoch, mean, std): data = dataset.make_one_shot_iterator().get_next() images = [] standardized = [] scans = sess.run(data) recons = sess.run(out, feed_dict={inp: scans}) for i, (scan, recon) in enumerate(zip(scans, recons)): scan_cross = np.array(scan[scan.shape[1] // 2, :, :]) recon_cross = np.array(recon[scan.shape[1] // 2, :, :]) standardized.append( np.hstack( (_normalize_image(scan_cross), _normalize_image(recon_cross)))) scan *= std scan += mean recon *= std recon += mean scan_cross = np.array(scan[scan.shape[1] // 2, :, :]) recon_cross = np.array(recon[scan.shape[1] // 2, :, :]) images.append( np.hstack( (_normalize_image(scan_cross), _normalize_image(recon_cross)))) # plt.figure() # plt.title(title) # _, axarr = plt.subplots(ncols=2) # axarr[0].title.set_text("Original") # axarr[0].imshow(scan_cross) # axarr[1].title.set_text("Reconstructed") # axarr[1].imshow(recon_cross) # plt.savefig(os.path.join(weights_path, "epoch{}_{}.png".format(epoch, i))) # plt.close('all') _tboard.log_images("{} reconstruction".format(title), images, step=epoch) _tboard.log_images("{} reconstruction (standardized)".format(title, ), standardized, step=epoch)
def _log_reconstruction_imgs(title, sess, dataset, inp, out, epoch, weights_path): data = dataset.make_one_shot_iterator().get_next() images = [] scans = sess.run(data) recons = sess.run(out, feed_dict={inp: scans}) for i, (scan, recon) in enumerate(zip(scans, recons)): scan_cross = scan[:, scan.shape[1] // 2, :] recon_cross = recon[:, scan.shape[1] // 2, :] plt.figure() plt.title(title) _, axarr = plt.subplots(ncols=2) axarr[0].title.set_text("Original") axarr[0].imshow(scan_cross) axarr[1].title.set_text("Reconstructed") axarr[1].imshow(recon_cross) plt.savefig( os.path.join(weights_path, "epoch{}_{}.png".format(epoch, i))) combined = np.hstack((scan_cross, recon_cross)) images.append(combined) _tboard.log_images("{} reconstruction".format(title), images, step=epoch)
def visualize(data, warpped, global_step, sid, opt, mode='both', name=''): def draw(img, pts, mask, color=None): res = img.copy() assert (pts.shape[0] == opt.max_matches) assert (mask.shape[0] == opt.max_matches) pts = (pts / 2 + .5) * (img.shape[:2])[::-1] # print('pts={}'.format(pts)) pts = pts.astype(np.int32) for i in range(pts.shape[0]): if not mask[i]: continue cv2.circle(res, tuple(pts[i]), 5, tuple(np.random.rand(3)) if color is None else color) return res prefix = util.train2show( torch.stack(list(map(lambda x: x[sid], data.prefix)), dim=0).data) unstable = util.train2show( torch.stack(list(map(lambda x: x[sid], data.unstable)), dim=0).data) target = util.train2show( torch.stack(list(map(lambda x: x[sid], data.target)), dim=0).data) warpped = util.train2show( torch.stack(list(map(lambda x: x[sid], warpped)), dim=0).data) diff = torch.abs(warpped - target) diff = to_gray(diff) diff = torch.stack([diff, diff, diff], dim=1) fm = list(map(lambda x: x.data, data.fm)) fm_mask = list(map(lambda x: x.data, data.fm_mask)) for i in range(len(fm)): pts = fm[i][sid].cpu().numpy() mask = fm_mask[i][sid].cpu().numpy() img = target[i].cpu().numpy().transpose([1, 2, 0]) # print(img.shape, pts.shape, mask.shape) # print('mask={}'.format(mask)) img = draw(img, pts[:, :2], mask) # stable target[i].copy_(torch.from_numpy(img.transpose([2, 0, 1]))) img = unstable[i].cpu().numpy().transpose([1, 2, 0]) img = draw(img, pts[:, 2:], mask) # unstable unstable[i].copy_(torch.from_numpy(img.transpose([2, 0, 1]))) vis = torch.cat((unstable, warpped, target, diff), dim=0) expr_dir = os.path.join(opt.checkpoints_dir, opt.name) prefix_grid = torchvision.utils.make_grid(prefix, nrow=1) vis_grid = torchvision.utils.make_grid(vis, nrow=vis.shape[0] // 4) if name != '': name += '-' if mode == 'both' or mode == 'save': torchvision.utils.save_image( prefix, os.path.join( expr_dir, name + 'prefix-{:0>4}-{:0>3}.png'.format(global_step, sid)), nrow=1) torchvision.utils.save_image( vis, os.path.join( expr_dir, name + 'input-output-target-{:0>4}-{:0>3}.png'.format( global_step, sid)), nrow=vis.shape[0] // 4) if name != '': name = name[:-1] + '/' if mode == 'both' or mode == 'log': tl.log_images(name + 'prefix/{}'.format(sid), [prefix_grid.cpu().numpy()], step=global_step) tl.log_images(name + 'input-output-target/{}'.format(sid), [vis_grid.cpu().numpy()], step=global_step) tl.log_images(name + 'diff/{}'.format(sid), diff[:, 0, ...].cpu().numpy(), step=global_step)
for path, objectid, pred_mask in zip(paths, objectids, preds_mask): logging.debug('Working on item "%s"' % path) # Write the probability. prob = pred_mask[0] # Take the first channel, which is the object. prob = np.uint8(prob * 255) prob = cv2.resize(prob, dsize=test_img_shape, interpolation=cv2.INTER_NEAREST) writer_prob.addImage(mask=prob, imagefile=path, width=test_img_shape[0], height=test_img_shape[1]) writer_prob.addObject({'objectid': objectid, 'imagefile': path}) if index % 10 == 0: log_images('test/gt/imgs', imgs[:1].cpu().data, step=index) log_images('test/pred/masks', torch.Tensor(255 - prob).unsqueeze(0), step=index) # Write the argmax class argmax_mask = np.argmax(pred_mask, axis=0) pred_mask = 255 - np.uint8(argmax_mask * 255) pred_mask = cv2.resize(pred_mask, dsize=test_img_shape, interpolation=cv2.INTER_NEAREST) writer_top.addImage(mask=pred_mask, imagefile=path) writer_top.addObject({'objectid': objectid, 'imagefile': path}) writer_prob.close() writer_top.close()
def train(args): # tensorboard run_name = "./runs/run-vae_batch_" + str(args.batch_size) \ + "_epochs_" + str(args.epochs) + "_" + args.log_message configure(run_name) # put on gpu if available device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") net = VAE().to(device) mnistmTrainSet = ds.mnistmTrainingDataset(text_file=args.dataset_list) mnistmTrainLoader = torch.utils.data.DataLoader(mnistmTrainSet, batch_size=args.batch_size, shuffle=True, num_workers=2) criterion = nn.BCELoss(size_average=False) optimizer = optim.Adam(net.parameters()) epoch = [0] # load prev model tio.load_model(model=net, optimizer=optimizer, epoch=epoch, path=run_name + '/ckpt') epoch = epoch[0] while epoch < args.epochs: for i, sample_batched in enumerate(mnistmTrainLoader, 0): input_batch = sample_batched['image'].float() input_batch = input_batch.to(device) input_batch = input_batch.view(-1, image_size) reconstruct, mu, log_var = net(input_batch) reconstruct_loss = criterion(reconstruct, input_batch) kl_loss = -0.5 * torch.sum(1 + log_var - mu.pow(2) - log_var.exp()) loss = reconstruct_loss + kl_loss optimizer.zero_grad() loss.backward() optimizer.step() # print statistics if i % 200 == 0: count = int( epoch * math.floor(len(mnistmTrainSet) / (args.batch_size * 200)) + (i / 200)) print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, loss.item())) log_value('loss', loss.item(), count) z = torch.randn(args.batch_size, z_dim).to(device) gen = net.decode(z).view(-1, 1, 28, 28) log_images('generated', gen[0].detach(), count) out, _, _ = net(input_batch) out = out.view(args.batch_size, 1, 28, 28) log_images('image', sample_batched['image'][0], count) log_images('recon', out[0].detach(), count) # save model tio.save_model(epoch=epoch, model=net, optimizer=optimizer, path=run_name + '/ckpt') epoch = epoch + 1
def val_rand(epoch): avg_loss = 0 avg_psnr = 0 frame_count = 0 epoch_start = time.time() end = time.time() average_meter = AverageMeter() modelME.module.eval() modelME.module.netM.eval() modelME.module.netE.eval() for batch in val_loader: image_target, depth_target, depth_mask = Variable(batch[0]), Variable( batch[1]), Variable(batch[2]) batch_size = depth_target.size(0) depth_input = depth_target.clone() image_height = image_target.size(2) image_width = image_target.size(3) # Corrupt the target image for i in range(0, batch_size): n_depth_mask = depth_mask[i, 0, :, :].sum().item() # Adjust the sampling rate based on the depth_mask sample_rate = opt.sample_rate / n_depth_mask * (image_height * image_width) corrupt_mask = np.random.binomial(1, (1 - sample_rate), (image_height, image_width)) corrupt_mask.astype(np.bool) corrupt_mask = torch.BoolTensor(corrupt_mask) depth_input[i, 0, :, :].masked_fill_(corrupt_mask, (0.0)) if torch.cuda.is_available(): image_target = image_target.cuda() depth_input = depth_input.cuda() depth_target = depth_target.cuda() depth_mask = depth_mask.cuda() rgb_sparse_d_input = torch.cat((image_target, depth_input), 1) torch.cuda.synchronize() data_time = time.time() - end # compute output end = time.time() with torch.no_grad(): depth_prediction = modelME.module.netE(rgb_sparse_d_input) torch.cuda.synchronize() gpu_time = time.time() - end # measure accuracy and record loss result = Result() result.evaluate(depth_prediction.data, depth_target.data) average_meter.update(result, gpu_time, data_time, depth_input.size(0)) end = time.time() for i in range(0, depth_input.shape[0]): loss_depth = criterion_depth(depth_prediction[[i]], depth_target[[i]], depth_mask[[i]]) loss_mse = criterion_mse(depth_prediction[[i]], depth_target[[i]], depth_mask[[i]]) psnr = 10 * log10(1 / loss_mse.data.item()) avg_loss += loss_depth.data.item() avg_psnr += psnr frame_count += 1 avg = average_meter.average() print('\n*\n' 'RMSE={average.rmse:.3f}\n' 'MAE={average.mae:.3f}\n' 'Delta1={average.delta1:.3f}\n' 'REL={average.absrel:.3f}\n' 'Lg10={average.lg10:.3f}\n' 't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time)) with open(val_rand_csv, 'a') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writerow({ 'mse': avg.mse, 'rmse': avg.rmse, 'absrel': avg.absrel, 'lg10': avg.lg10, 'mae': avg.mae, 'delta1': avg.delta1, 'delta2': avg.delta2, 'delta3': avg.delta3, 'data_time': avg.data_time, 'gpu_time': avg.gpu_time }) epoch_end = time.time() print( "===> Epoch {} Random Validation: Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Time: {:.4f}" .format(epoch, avg_loss / frame_count, avg_psnr / frame_count, (epoch_end - epoch_start))) log_value('val_loss_depth_rand', avg_loss / frame_count, epoch) log_value('val_psnr_rand', avg_psnr / frame_count, epoch) log_images('depth_input_rand', reshape_4D_array(depth_input[[0], :, :, :].data.cpu().numpy(), 1), step=1) log_images('depth_prediction_rand', reshape_4D_array( depth_prediction[[0], :, :, :].data.cpu().numpy(), 1), step=1) log_images('depth_target_rand', reshape_4D_array(depth_target[[0], :, :, :].data.cpu().numpy(), 1), step=1) log_images('image_target_rand', reshape_4D_array(image_target[[0], :, :, :].data.cpu().numpy(), 1), step=1)
def val(epoch): avg_loss = 0 avg_loss_depth = 0 avg_loss_pos = 0 avg_loss_temperature = 0 avg_psnr = 0 avg_sparsity = 0 avg_loss_slic = 0.0 avg_loss_color = 0.0 avg_loss_pos = 0.0 frame_count = 0 epoch_start = time.time() end = time.time() average_meter = AverageMeter() modelME.module.eval() modelME.module.netM.eval() modelME.module.netE.eval() for batch in val_loader: image_target, depth_target, depth_mask = Variable(batch[0]), Variable( batch[1]), Variable(batch[2]) image_target_lab = rgb2Lab_torch( image_target.cuda()) # image in lab color space image_target_labxy_feat_tensor = build_LABXY_feat( image_target_lab, train_XY_feat_stack) # B* (3+2 )* H * W depth_input = depth_target.clone() if torch.cuda.is_available(): image_target = image_target.cuda() depth_input = depth_input.cuda() depth_target = depth_target.cuda() depth_mask = depth_mask.cuda() image_target_labxy_feat_tensor = image_target_labxy_feat_tensor.cuda( ) torch.cuda.synchronize() data_time = time.time() - end # compute output end = time.time() with torch.no_grad(): depth_recon, corrupt_mask_soft, corrupt_mask_binary, pooled_xy_tensor, reconstr_xy_tensor, curr_spixl_map, prob = modelME( image_target, depth_input, False) torch.cuda.synchronize() mask_sparsity = corrupt_mask_binary.sum() / ( corrupt_mask_binary.shape[0] * corrupt_mask_binary.shape[1] * corrupt_mask_binary.shape[2] * corrupt_mask_binary.shape[3]) # Generate the corrupted depth image depth_input = corrupt_mask_binary * depth_input rgb_sparse_d_input = torch.cat((image_target, depth_input), 1) # white input with torch.no_grad(): restored_depth = modelME.module.netE(rgb_sparse_d_input) torch.cuda.synchronize() gpu_time = time.time() - end # measure accuracy and record loss result = Result() result.evaluate(restored_depth.data, depth_target.data) average_meter.update(result, gpu_time, data_time, depth_input.size(0)) end = time.time() for i in range(0, depth_input.shape[0]): loss_depth = criterion_depth(restored_depth[[i]], depth_target[[i]], depth_mask[[i]]) loss_mse = criterion_mse(restored_depth[[i]], depth_target[[i]], depth_mask[[i]]) psnr = 10 * log10(1 / loss_mse.data.item()) loss_slic, loss_color, loss_pos = compute_color_pos_loss( prob[[i]], image_target_labxy_feat_tensor[[i]], pos_weight=opt.pos_weight, kernel_size=opt.downsize) loss_temperature = modelME.module.netM.temperature**2 loss = loss_depth + opt.slic_weight * loss_slic + opt.proj_weight * loss_temperature avg_loss += loss.data.item() avg_loss_depth += loss_depth.data.item() avg_loss_pos += loss_pos.data.item() avg_loss_temperature += loss_temperature.data.item() avg_psnr += psnr avg_sparsity += mask_sparsity avg_loss_slic += loss_slic.data.item() avg_loss_color += loss_color.data.item() avg_loss_pos += loss_pos.data.item() frame_count += 1 avg = average_meter.average() print('\n*\n' 'RMSE={average.rmse:.3f}\n' 'MAE={average.mae:.3f}\n' 'Delta1={average.delta1:.3f}\n' 'REL={average.absrel:.3f}\n' 'Lg10={average.lg10:.3f}\n' 't_GPU={time:.3f}\n'.format(average=avg, time=avg.gpu_time)) with open(val_csv, 'a') as csvfile: writer = csv.DictWriter(csvfile, fieldnames=fieldnames) writer.writerow({ 'mse': avg.mse, 'rmse': avg.rmse, 'absrel': avg.absrel, 'lg10': avg.lg10, 'mae': avg.mae, 'delta1': avg.delta1, 'delta2': avg.delta2, 'delta3': avg.delta3, 'data_time': avg.data_time, 'gpu_time': avg.gpu_time }) epoch_end = time.time() print( "===> Epoch {} Validation: Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Mask Sparsity: {:.4f}, Time: {:.4f}" .format(epoch, avg_loss / frame_count, avg_psnr / frame_count, avg_sparsity / frame_count, (epoch_end - epoch_start))) log_value('val_loss', avg_loss / frame_count, epoch) log_value('val_loss_depth', avg_loss_depth / frame_count, epoch) log_value('val_loss_pos', avg_loss_pos / frame_count, epoch) log_value('val_loss_temperature', avg_loss_temperature / frame_count, epoch) log_value('val_psnr', avg_psnr / frame_count, epoch) log_value('val_sparsity', avg_sparsity / frame_count, epoch) log_value('val_loss_slic', avg_loss_slic / frame_count, epoch) log_value('val_loss_color', avg_loss_color / frame_count, epoch) log_value('val_loss_pos', avg_loss_pos / frame_count, epoch) # Draw the last image result spixl_map = update_spixl_map(val_spixelID[[-1], :, :, :], prob[[-1], :, :, :]) # 1x1x240x960 ori_sz_spixel_map = F.interpolate( spixl_map.type(torch.float), size=(opt.input_img_height, opt.input_img_width), mode='nearest').type(torch.int) # 1x1x240x960 spix_index_np = ori_sz_spixel_map.squeeze().detach().cpu().numpy( ).transpose(0, 1) #240x960 spix_index_np = spix_index_np.astype(np.int64) #240x960, 1% here image_rgb_np = image_target[[-1], :, :, :].squeeze().clamp( 0, 1).detach().cpu().numpy().transpose(1, 2, 0) spixel_bd_image = mark_boundaries(image_rgb_np, spix_index_np.astype(int), color=(0, 1, 1)) spixel_viz = spixel_bd_image.astype(np.float32).transpose(2, 0, 1) spixel_viz = np.expand_dims(spixel_viz, axis=0) log_images('image_target_spixel', reshape_4D_array(spixel_viz, 1), step=1) log_images('depth_input', reshape_4D_array(depth_input[[-1], :, :, :].data.cpu().numpy(), 1), step=1) log_images('depth_prediction', reshape_4D_array( restored_depth[[-1], :, :, :].data.cpu().numpy(), 1), step=1) log_images('depth_target', reshape_4D_array(depth_target[[-1], :, :, :].data.cpu().numpy(), 1), step=1) log_images('image_target', reshape_4D_array(image_target[[-1], :, :, :].data.cpu().numpy(), 1), step=1) log_images('corrupt_mask_binary', reshape_4D_array( corrupt_mask_binary[[-1], :, :, :].data.cpu().numpy(), 1), step=1) log_images('corrupt_mask_soft', reshape_4D_array( corrupt_mask_soft[[-1], :, :, :].data.cpu().numpy(), 1), step=1) global LOSS_best if avg_loss_depth < LOSS_best: LOSS_best = avg_loss model_out_path = opt.path_to_save + "/model_best.pth" torch.save(modelME.module.state_dict(), model_out_path) print("Checkpoint saved to {}".format(model_out_path))
get_data(DATA_DIR + '/goog.csv') dates = np.reshape(dates, (len(dates), 1)) svm = SVR(kernel=kernel, C=C, degree=degree, gamma=gamma) svm.fit(dates, prices) predictions = svm.predict(dates) (rmse, mae, r2) = eval_metrics(prices, predictions) log_metrics('RMSE', rmse) log_metrics('MAE', mae) log_metrics('R2', r2) plt.plot(dates, prices, color='black', label="Data", marker='*') plt.plot(dates, predictions, color='red', label="Predictions", marker='o') plt.xlabel('Date') plt.ylabel('Price') plt.title('SVM predictions with ' + kernel + ' kernel') plt.legend() plt.savefig('svm.png') log_value('RMSE', rmse) log_value('MAE', mae) log_value('R2', r2) filename = MODEL_DIR + '/model.joblib' joblib.dump(svm, filename) img = cv2.imread('svm.png') log_histogram('Stock Prices', prices, step=1) log_images('Stock Predictions Graph', [img])
if __name__ == "__main__": print ("MODEL_DIR:{}, DATA_DIR:{}".format(MODEL_DIR,DATA_DIR)) get_data(DATA_DIR +'/goog.csv') dates = np.reshape(dates,(len(dates), 1)) svm = SVR(kernel= kernel, C= C, degree= degree, gamma=gamma) svm.fit(dates, prices) predictions = svm.predict(dates) (rmse, mae, r2) = eval_metrics(prices, predictions) plt.scatter(dates, prices, color= 'black', label= 'Data') plt.plot(dates,predictions, color= 'red', label= 'SVM model') plt.xlabel('Date') plt.ylabel('Price') plt.title('SVM with '+kernel+ ' kernel') plt.legend() plt.savefig('svm.png') log_value('RMSE', rmse) log_value('MAE', mae) log_value('R2', r2) img = cv2.imread('svm.png') log_histogram('Stock Prices', prices, step=1) log_images('Stock Prediction Model',[img])
def val(epoch): avg_psnr = 0 avg_mse = 0 avg_sparsity = 0 modelME.eval() modelME.netM.eval() modelME.netE.eval() for batch in val_loader: target = batch image = target.clone() image_clone = image.clone() mean_image = torch.zeros(image.shape[0], image.shape[1], image.shape[2], image.shape[3]) mean_image[:,0,:,:] = 0.5 mean_image[:,1,:,:] = 0.5 mean_image[:,2,:,:] = 0.5 std_image = torch.zeros(image.shape[0], image.shape[1], image.shape[2], image.shape[3]) std_image[:,0,:,:] = 0.5 std_image[:,1,:,:] = 0.5 std_image[:,2,:,:] = 0.5 if torch.cuda.is_available(): image = image.cuda() image_clone = image_clone.cuda() target = target.cuda() mean_image = mean_image.cuda() std_image = std_image.cuda() # Generate the corruption mask and reconstructed image corrupt_mask_conti, image_recon = modelME(image) corrupt_mask = corrupt_mask_conti.bernoulli() # Binarize the corruption mask using Bernoulli distribution, then feed into modelE mask_sparsity = corrupt_mask.sum() / (corrupt_mask.shape[0] * corrupt_mask.shape[1] * corrupt_mask.shape[2] * corrupt_mask.shape[3]) corrupt_mask = corrupt_mask.expand(corrupt_mask.shape[0], 3, corrupt_mask.shape[2], corrupt_mask.shape[3]) # Generate the corrupted image mask_image = corrupt_mask * image_clone restored_image = modelME.netE(mask_image) mse = criterion((restored_image*std_image)+mean_image, (target*std_image)+mean_image) psnr = 10 * log10(1 / mse.data[0]) avg_psnr += psnr avg_mse += mse.data[0] avg_sparsity += mask_sparsity print("===> Epoch {} Validation: Avg. Loss: {:.4f}, Avg.PSNR: {:.4f} dB, Mask Sparsity: {:.4f}".format(epoch, avg_mse / len(val_loader), avg_psnr / len(val_loader), avg_sparsity / len(val_loader))) log_value('val_loss', avg_mse / len(val_loader), epoch) log_value('val_psnr', avg_psnr / len(val_loader), epoch) log_value('val_sparsity', avg_sparsity / len(val_loader), epoch) corrupt_mask_conti = corrupt_mask_conti.expand(corrupt_mask_conti.shape[0], 3, corrupt_mask_conti.shape[2], corrupt_mask_conti.shape[3]) log_images('original_image', reshape_4D_array((image*std_image+mean_image).cpu().numpy(), 10), step=1) log_images('conti_mask', reshape_4D_array(corrupt_mask_conti.data.cpu().numpy(), 10), step=1) log_images('binar_mask', reshape_4D_array(corrupt_mask.data.cpu().numpy(), 10), step=1) log_images('restored_image', reshape_4D_array((restored_image*std_image+mean_image).data.cpu().numpy(), 10), step=1) global PSNR_best if avg_psnr > PSNR_best: PSNR_best = avg_psnr model_out_path = "epochs_NetME/" + "model_best.pth" torch.save(modelME.state_dict(), model_out_path) print("Checkpoint saved to {}".format(model_out_path))
def val(epoch): avg_total_loss = 0.0 avg_sem_loss = 0.0 avg_pos_loss = 0.0 frame_count = 0 epoch_start = time.time() model.module.eval() for batch in val_loader: image_rgb, depth_target, depth_mask = Variable(batch[0]), Variable( batch[1]), Variable(batch[2]) image_lab = rgb2Lab_torch(image_rgb.cuda()) # image in lab color space LABXY_feat_tensor = build_LABXY_feat(image_lab, train_XY_feat_stack) # B* (3+2) if torch.cuda.is_available(): image_rgb = image_rgb.cuda() LABXY_feat_tensor = LABXY_feat_tensor.cuda() torch.cuda.synchronize() with torch.no_grad(): output = model(image_rgb) # white output torch.cuda.synchronize() slic_loss, loss_sem, loss_pos = compute_color_pos_loss( output, LABXY_feat_tensor, pos_weight=opt.pos_weight, kernel_size=opt.downsize) avg_total_loss += slic_loss.data.item() avg_sem_loss += loss_sem.data.item() avg_pos_loss += loss_pos.data.item() epoch_end = time.time() print( "===> Epoch {} Validation: Avg. total_loss: {:.4f}, Avg. sem_loss: {:.4f}, Avg. epoch_pos_loss: {:.4f}, Time: {:.4f}" .format(epoch, avg_total_loss / len(val_loader), avg_sem_loss / len(val_loader), avg_pos_loss / len(val_loader), (epoch_end - epoch_start))) log_value('val_total_loss', avg_total_loss / len(val_loader), epoch) log_value('val_sem_loss', avg_sem_loss / len(val_loader), epoch) log_value('val_pos_loss', avg_pos_loss / len(val_loader), epoch) # Draw the last image result spixl_map = update_spixl_map(val_spixelID[[-1], :, :, :], output[[-1], :, :, :]) # 1x1x240x960 ori_sz_spixel_map = F.interpolate( spixl_map.type(torch.float), size=(opt.input_img_height, opt.input_img_width), mode='nearest').type(torch.int) # 1x1x240x960 spix_index_np = ori_sz_spixel_map.squeeze().detach().cpu().numpy( ).transpose(0, 1) #240x960 spix_index_np = spix_index_np.astype(np.int64) #240x960, 1% here image_rgb_np = image_rgb[[-1], :, :, :].squeeze().clamp( 0, 1).detach().cpu().numpy().transpose(1, 2, 0) spixel_bd_image = mark_boundaries(image_rgb_np, spix_index_np.astype(int), color=(0, 1, 1)) spixel_viz = spixel_bd_image.astype(np.float32).transpose(2, 0, 1) spixel_viz = np.expand_dims(spixel_viz, axis=0) image_rgb_np_viz = image_rgb_np.astype(np.float32).transpose(2, 0, 1) image_rgb_np_viz = np.expand_dims(image_rgb_np_viz, axis=0) log_images('spixel', reshape_4D_array(spixel_viz, 1), step=1) log_images('image_rgb', reshape_4D_array(image_rgb_np_viz, 1), step=1) global LOSS_best if avg_total_loss < LOSS_best: LOSS_best = avg_total_loss model_out_path = opt.path_to_save + "/model_best.pth".format(epoch) torch.save(model.module.state_dict(), model_out_path) print("Checkpoint saved to {}".format(model_out_path))
def log_images(harn, key, value, n_iter): if False: print('{}={} @ {}'.format(key, value, n_iter)) if tensorboard_logger: tensorboard_logger.log_images(key, value, n_iter)
def train(args): # tensorboard run_name = "./runs/run-GAN_batch_" + str(args.batch_size) \ + "_epochs_" + str(args.epochs) + "_" + args.log_message configure(run_name) gen = generator() desc = descrimanator() mnistmTrainSet = ds.mnistmTrainingDataset(text_file=args.dataset_list) mnistmTrainLoader = torch.utils.data.DataLoader(mnistmTrainSet, batch_size=args.batch_size, shuffle=True, num_workers=2) # put on gpu if available device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") loss = torch gen.to(device) desc.to(device) gen_optimizer = optim.Adam(gen.parameters()) desc_optimizer = optim.SGD(desc.parameters(), lr=0.0001, momentum=0.9) criterion = nn.BCELoss() epoch = [0] #load prev model tio.load_model(model=gen, optimizer=gen_optimizer, epoch=epoch, path=run_name + '/ckpt_gan') tio.load_model(model=desc, optimizer=desc_optimizer, epoch=epoch, path=run_name + '/ckpt_desc') epoch = epoch[0] while epoch < args.epochs: for i, sample_batched in enumerate(mnistmTrainLoader, 0): input_batch = sample_batched['image'].float() input_batch = 2 * (input_batch - 0.5) input_batch = input_batch.to(device) # Create the labels which are later used as input for the BCE loss # ones_labels = torch.ones(input_batch.shape[0], 1, 1, 1).to(device) real_labels = 0.99 * torch.ones(input_batch.shape[0], 1, 1, 1).to(device) fake_labels = 0.01 * torch.ones(input_batch.shape[0], 1, 1, 1).to(device) # ================================================================== # # Train the discriminator # # ================================================================== # real_desc = desc(input_batch) desc_loss_real = criterion(real_desc, real_labels) noise = torch.randn(input_batch.shape[0], 1, 1, 100).to(device) gen_imgs = gen(noise) gen_desc = desc(gen_imgs) desc_loss_fake = criterion(gen_desc, fake_labels) desc_loss = desc_loss_fake + desc_loss_real gen_optimizer.zero_grad() desc_optimizer.zero_grad() desc_loss.backward() desc_optimizer.step() # ================================================================== # # Train the generator # # ================================================================== # noise = torch.empty(input_batch.shape[0], 1, 1, 100).uniform_().to(device) gen_imgs = gen(noise) gen_desc = desc(gen_imgs) gen_loss = criterion(gen_desc, real_labels) gen_optimizer.zero_grad() desc_optimizer.zero_grad() gen_loss.backward() gen_optimizer.step() if (i + 1) % 200 == 0: count = int( epoch * math.floor(len(mnistmTrainSet) / (args.batch_size * 200)) + (i / 200)) print( 'Epoch [{}], Step [{}], desc_loss: {:.4f}, gen_loss: {:.4f}, D(x): {:.2f}, D(G(z)): {:.2f}' .format(epoch, i + 1, desc_loss.item(), gen_loss.item(), real_desc.mean().item(), gen_desc.mean().item())) log_value("Gen Loss", gen_loss.item(), count) log_value("Desc Loss", desc_loss.item(), count) log_value("D(x)", real_desc.mean().item(), count) log_value("D(G(z))", gen_desc.mean().item(), count) for i in range(input_batch.shape[0]): log_images("generated", gen_imgs[i].detach(), count) # train generator #save model tio.save_model(model=gen, optimizer=gen_optimizer, epoch=epoch, path=run_name + '/ckpt_gan') tio.save_model(model=desc, optimizer=desc_optimizer, epoch=epoch, path=run_name + '/ckpt_desc') epoch = epoch + 1
pred_tgt_masks2, _ = model_f2(features) c_loss_mask = criterion_mask( pred_tgt_masks1[tgt_use_masks, :, :, :], tgt_masks[tgt_use_masks, :, :]) c_loss_mask += criterion_mask( pred_tgt_masks2[tgt_use_masks, :, :, :], tgt_masks[tgt_use_masks, :, :]) c_loss_mask /= 2. c_loss_mask[torch.isnan(c_loss_mask)] = 0 c_loss = c_loss_mask c_loss.backward() tflogger.acc_value( 'train/loss/mask_tgt', c_loss_mask / tgt_use_masks.sum().type(torch.float)) log_images('train/gt/tgt_masks', tgt_masks[:1].cpu().data, step=log_counter) optimizer_f.step() optimizer_g.step() if (log_counter + 1) % args.freq_log == 0: print 'Epoch %d/%d, batch %d/%d' % \ (epoch, args.epochs, ibatch, len(train_loader)), tflogger.get_mean_values() log_images('train/gt/src_imgs', src_imgs[:1].cpu().data, step=log_counter) log_images('train/gt/tgt_imgs', tgt_imgs[:1].cpu().data, step=log_counter) log_images('train/gt/src_masks',
def train(self, epoch, train_loader, metrics): ''' Trains the model for a single epoch ''' # train_size = int(0.9 * len(train_loader.dataset.train_data) / self.config.batch_size) loss_sum = 0.0 N = 0 # print('\33[1m==> Training epoch # {}\033[0m'.format(str(epoch))) classified = [] misclassified = [] self.model.train() budget_exceeded = False metric_results = [0] * len(metrics) start_time = time.time() for step, (data, targets) in enumerate(train_loader): # import matplotlib.pyplot as plt # img = plt.imshow(data.numpy()[0,1,:]) # plt.show() # images += list(data.numpy()) # print('Data:', data.size(), ' - Label:', targets.size()) data = data.to(self.device) targets = targets.to(self.device) data, criterion_kwargs = self.loss_computation.prepare_data(data, targets) batch_size = data.size(0) outputs = self.model(data) loss_func = self.loss_computation.criterion(**criterion_kwargs) loss = loss_func(self.criterion, outputs) self.optimizer.zero_grad() loss.backward() self.optimizer.step() # print('Train:', ' '.join(str(outputs).split('\n')[0:2])) if self.images_plot_count > 0: with torch.no_grad(): _, pred = outputs.topk(1, 1, True, True) pred = pred.t() correct = pred.eq(targets.view(1, -1).expand_as(pred)).cpu().numpy()[0] data = data.cpu().numpy() classified += list(data[correct.astype(bool)]) misclassified += list(data[(1-correct).astype(bool)]) if len(classified) > self.images_plot_count: classified = random.sample(classified, self.images_plot_count) if len(misclassified) > self.images_plot_count: misclassified = random.sample(misclassified, self.images_plot_count) # self.scheduler.cumulative_time += delta_time # self.scheduler.last_step = self.scheduler.cumulative_time - delta_time - 1e-10 tmp = time.time() with torch.no_grad(): for i, metric in enumerate(metrics): metric_results[i] += self.loss_computation.evaluate(metric, outputs, **criterion_kwargs) * batch_size loss_sum += loss.item() * batch_size N += batch_size #print('Update', (metric_results[0] / N), 'loss', (loss_sum / N), 'lr', self.optimizer.param_groups[0]['lr']) if self.budget_type == 'time' and self.cumulative_time + (time.time() - start_time) >= self.budget: # print(' * Stopping at Epoch: [%d][%d/%d] for a budget of %.3f s' % (epoch, step + 1, train_size, self.config.budget)) budget_exceeded = True break if N==0: # Fixes a bug during initialization N=1 if self.images_plot_count > 0: import tensorboard_logger as tl tl.log_images('Train_Classified/Image', classified, step=epoch) tl.log_images('Train_Misclassified/Image', misclassified, step=epoch) if self.checkpoint_path and self.scheduler.snapshot_before_restart and self.scheduler.needs_checkpoint(): self.latest_checkpoint = save_checkpoint(self.checkpoint_path, self.config_id, self.budget, self.model, self.optimizer, self.scheduler) try: self.scheduler.step(epoch=epoch) except: self.scheduler.step(metrics=loss_sum / N, epoch=epoch) self.cumulative_time += (time.time() - start_time) #print('LR', self.optimizer.param_groups[0]['lr'], 'Update', (metric_results[0] / N), 'loss', (loss_sum / N)) return [res / N for res in metric_results], loss_sum / N, budget_exceeded
pred_masks, _ = model_f1(features) loss_mask = criterion_mask(pred_masks, gt_masks) loss = loss_mask loss.backward() tflogger.acc_value('train/loss/total', loss / args.batch_size) tflogger.acc_value('train/loss/mask', loss_mask / args.batch_size) optimizer_g.step() optimizer_f.step() if (log_counter + 1) % args.freq_log == 0: print 'Epoch %d/%d, batch %d/%d' % \ (epoch, args.epochs, ibatch, len(train_loader)), tflogger.get_mean_values() log_images('train/gt/imgs', imgs[:1].cpu().data, step=log_counter) log_images('train/gt/masks', gt_masks[:1].cpu().data, step=log_counter) pred_masks_asimage = torch.softmax(pred_masks, dim=1)[:1, 1].cpu().data #print ('pred_masks_asimage', pred_masks_asimage.min(), pred_masks_asimage.max(), pred_masks_asimage.dtype) log_images('train/pred/masks', pred_masks_asimage, step=log_counter) tflogger.flush(step=log_counter - 1) if args.max_iter is not None and istep >= args.max_iter: break log_value('lr', args.lr, epoch)