default="./logs/81norm_100p/net.pth", help="Existing model") parser.add_argument("--batchSize", type=int, default=100, help="batch size") parser.add_argument("--noiselevel", type=int, default=15, help="noise") parser.add_argument("--num_of_layers", type=int, default=9, help="layers") parser.add_argument("--outf", type=str, default="", help="output director path") parser.add_argument("--patchSize", type=int, default=50, help="patch size") args = parser.parse_args() model = DnCNN() model.load_state_dict(torch.load(args.model)) model.eval() model.to('cuda') branch = get_all_histograms("./test.root") model.to('cpu') data_means = np.zeros(50) noisy_means = np.zeros(50) output_means = np.zeros(50) n_d_ratio = np.zeros(50) o_d_ratio = np.zeros(50) for image in range(50): data = get_bin_weights(branch, image).copy() data_mean = np.mean(data) stdevs = np.std(data)
training_set = DnCnnDataset(trainRootPath, noise_stddev, training=True) training_generator = DataLoader(training_set, batch_size=batch_size, shuffle=True, num_workers=4,drop_last=True) eval_set = DnCnnDataset(evalSetPath, noise_stddev, training=False, evaluate=False) eval_set_generator = DataLoader(eval_set,batch_size=batch_size, shuffle=False, drop_last=True) DnCnn_net = DnCNN(channels=batch_size, layers_num=10) #transfer to GPU - cuda/to.device criterion = nn.MSELoss() # Move to GPU model = DnCnn_net.to(device) optimizer = optim.Adam(model.parameters(), lr=lr) optim_scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [int(0.1*epochs_num), int(0.9*epochs_num)], gamma=0.1) psnr_values = [] for epoch in range(epochs_num): running_loss = 0.0 optim_scheduler.step() for i, data in enumerate(training_generator): image, noise, noised_image = data image, noise, noised_image = image.to(device).unsqueeze(0),\ noise.to(device).unsqueeze(0),noised_image.unsqueeze(0).to(device)
def main(): # creat_readme() # choose cpu or gpu if torch.cuda.is_available(): args.device = torch.device('cuda') else: args.device = torch.device('cpu') print('Loading Dataset--') dataset_train = RootDataset(root_file=args.trainfile, sigma=args.sigma) loader_train = DataLoader(dataset=dataset_train, batch_size=args.batchSize) dataset_val = RootDataset(root_file=args.valfile, sigma=args.sigma) val_train = DataLoader(dataset=dataset_val, batch_size=args.batchSize) # Build model model = DnCNN(channels=1, num_of_layers=args.num_of_layers, ker_size=args.kernelSize, o_k_size=args.outKerSize).to(device=args.device) if (args.model == None): model.apply(init_weights) print("Creating new model") else: print("Loading model from file" + args.model) model.load_state_dict(torch.load(args.model)) model.eval() # Loss function criterion = PatchLoss() criterion.to(device=args.device) #Optimizer MyOptim = optim.Adam(model.parameters(), lr=args.lr) MyScheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer=MyOptim, factor=0.1, patience=10, verbose=True) # training and validation step = 0 training_losses = np.zeros(args.epochs) validation_losses = np.zeros(args.epochs) for epoch in range(args.epochs): print("Epoch #" + str(epoch)) # training train_loss = 0 for i, data in enumerate(loader_train, 0): model.train() model.zero_grad() MyOptim.zero_grad() truth, noise = data noise = noise.unsqueeze(1) output = model(noise.float().to(args.device), args.outKerSize) batch_loss = criterion( output.squeeze(1).to(args.device), truth.to(args.device), args.patchSize).to(args.device) batch_loss.backward() MyOptim.step() model.eval() train_loss += batch_loss.item() training_losses[epoch] = train_loss print("Training Loss: " + str(train_loss)) val_loss = 0 for i, data in enumerate(val_train, 0): val_truth, val_noise = data val_output = model( val_noise.unsqueeze(1).float().to(args.device), args.outKerSize) output_loss = criterion( val_output.squeeze(1).to(args.device), val_truth.to(args.device), args.patchSize).to(args.device) val_loss += output_loss.item() MyScheduler.step(torch.tensor([val_loss])) validation_losses[epoch] = val_loss print("Validation Loss: " + str(val_loss)) # save the model model.eval() torch.save(model.state_dict(), os.path.join(args.outf, 'net.pth')) training = plt.plot(training_losses, label='Training') validation = plt.plot(validation_losses, label='Validation') plt.legend() plt.savefig(args.outf + "/lossplt.png") branch = get_all_histograms("./test.root") model.to('cpu') for image in range(10): data = get_bin_weights(branch, image).copy() np.savetxt(args.outf + '/truth#' + str(image) + '.txt', data) means = np.mean(data) stdevs = np.std(data) noisy = add_noise(data, args.sigma).copy() np.savetxt(args.outf + '/noisy#' + str(image) + '.txt', noisy) data_norm = (data - means) / stdevs np.savetxt(args.outf + '/truth_norm#' + str(image) + '.txt', data_norm) noisy_norm = (noisy - means) / stdevs np.savetxt(args.outf + '/noisy_norm#' + str(image) + '.txt', noisy_norm) data_norm = torch.from_numpy(data_norm) noisy_norm = torch.from_numpy(noisy_norm) noisy_norm = noisy_norm.unsqueeze(0) noisy_norm = noisy_norm.unsqueeze(1) output_norm = model( noisy_norm.float(), args.outKerSize).squeeze(0).squeeze(0).detach().numpy() np.savetxt(args.outf + '/output_norm#' + str(image) + '.txt', output_norm) output = (output_norm * stdevs) + means np.savetxt(args.outf + '/output#' + str(image) + '.txt', output) truth = data.numpy() noisy = noisy.numpy() diff = output - truth noisy_diff = noisy - truth np.savetxt(args.outf + '/diff#' + str(image) + '.txt', diff) model.to('cuda')