def train_model_residual_lowlight_rdn(): device = DEVICE #准备数据 train_set = HsiCubicTrainDataset('./data/train_lowlight_patchsize32/') #print('trainset32 training example:', len(train_set32)) #train_set = HsiCubicTrainDataset('./data/train_lowlight/') #train_set_64 = HsiCubicTrainDataset('./data/train_lowlight_patchsize64/') #train_set_list = [train_set32, train_set_64] #train_set = ConcatDataset(train_set_list) #里面的样本大小必须是一致的,否则会连接失败 print('total training example:', len(train_set)) train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=True) #加载测试label数据 mat_src_path = './data/test_lowlight/origin/soup_bigcorn_orange_1ms.mat' test_label_hsi = scio.loadmat(mat_src_path)['label'] #加载测试数据 batch_size = 1 #test_data_dir = './data/test_lowlight/cuk12/' test_data_dir = './data/test_lowlight/cubic/' test_set = HsiCubicLowlightTestDataset(test_data_dir) test_dataloader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=False) batch_size, channel, width, height = next(iter(test_dataloader))[0].shape band_num = len(test_dataloader) denoised_hsi = np.zeros((width, height, band_num)) save_model_path = './checkpoints/hsirnd_cosine' if not os.path.exists(save_model_path): os.mkdir(save_model_path) #创建模型 net = HSIRDN(K) init_params(net) net = nn.DataParallel(net).to(device) #net = net.to(device) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE) #scheduler = MultiStepLR(hsid_optimizer, milestones=[200,400], gamma=0.5) scheduler = CosineAnnealingLR(hsid_optimizer, T_max=600) #定义loss 函数 #criterion = nn.MSELoss() is_resume = RESUME #唤醒训练 if is_resume: path_chk_rest = dir_utils.get_last_path(save_model_path, 'model_latest.pth') model_utils.load_checkpoint(net, path_chk_rest) start_epoch = model_utils.load_start_epoch(path_chk_rest) + 1 model_utils.load_optim(hsid_optimizer, path_chk_rest) for i in range(1, start_epoch): scheduler.step() new_lr = scheduler.get_lr()[0] print( '------------------------------------------------------------------------------' ) print("==> Resuming Training with learning rate:", new_lr) print( '------------------------------------------------------------------------------' ) global tb_writer tb_writer = get_summary_writer(log_dir='logs') gen_epoch_loss_list = [] cur_step = 0 first_batch = next(iter(train_loader)) best_psnr = 0 best_epoch = 0 best_iter = 0 if not is_resume: start_epoch = 1 num_epoch = 600 for epoch in range(start_epoch, num_epoch + 1): epoch_start_time = time.time() scheduler.step() print('epoch = ', epoch, 'lr={:.6f}'.format(scheduler.get_lr()[0])) print(scheduler.get_lr()) gen_epoch_loss = 0 net.train() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for batch_idx, (noisy, cubic, label) in enumerate(train_loader): #print('batch_idx=', batch_idx) noisy = noisy.to(device) label = label.to(device) cubic = cubic.to(device) hsid_optimizer.zero_grad() #denoised_img = net(noisy, cubic) #loss = loss_fuction(denoised_img, label) residual = net(noisy, cubic) alpha = 0.8 loss = recon_criterion(residual, label - noisy) #loss = alpha*recon_criterion(residual, label-noisy) + (1-alpha)*loss_function_mse(residual, label-noisy) #loss = recon_criterion(residual, label-noisy) loss.backward() # calcu gradient hsid_optimizer.step() # update parameter gen_epoch_loss += loss.item() if cur_step % display_step == 0: if cur_step > 0: print( f"Epoch {epoch}: Step {cur_step}: Batch_idx {batch_idx}: MSE loss: {loss.item()}" ) else: print("Pretrained initial state") tb_writer.add_scalar("MSE loss", loss.item(), cur_step) #step ++,每一次循环,每一个batch的处理,叫做一个step cur_step += 1 gen_epoch_loss_list.append(gen_epoch_loss) tb_writer.add_scalar("mse epoch loss", gen_epoch_loss, epoch) #scheduler.step() #print("Decaying learning rate to %g" % scheduler.get_last_lr()[0]) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"{save_model_path}/hsid_rdn_4rdb_conise_l1_loss_600epoch_patchsize32_{epoch}.pth" ) #测试代码 net.eval() psnr_list = [] for batch_idx, (noisy_test, cubic_test, label_test) in enumerate(test_dataloader): noisy_test = noisy_test.type(torch.FloatTensor) label_test = label_test.type(torch.FloatTensor) cubic_test = cubic_test.type(torch.FloatTensor) noisy_test = noisy_test.to(DEVICE) label_test = label_test.to(DEVICE) cubic_test = cubic_test.to(DEVICE) with torch.no_grad(): residual = net(noisy_test, cubic_test) denoised_band = noisy_test + residual denoised_band_numpy = denoised_band.cpu().numpy().astype( np.float32) denoised_band_numpy = np.squeeze(denoised_band_numpy) denoised_hsi[:, :, batch_idx] = denoised_band_numpy if batch_idx == 49: residual_squeezed = torch.squeeze(residual, axis=0) denoised_band_squeezed = torch.squeeze(denoised_band, axis=0) label_test_squeezed = torch.squeeze(label_test, axis=0) noisy_test_squeezed = torch.squeeze(noisy_test, axis=0) tb_writer.add_image(f"images/{epoch}_restored", denoised_band_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_residual", residual_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_label", label_test_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_noisy", noisy_test_squeezed, 1, dataformats='CHW') test_label_current_band = test_label_hsi[:, :, batch_idx] psnr = PSNR(denoised_band_numpy, test_label_current_band) psnr_list.append(psnr) mpsnr = np.mean(psnr_list) denoised_hsi_trans = denoised_hsi.transpose(2, 0, 1) test_label_hsi_trans = test_label_hsi.transpose(2, 0, 1) mssim = SSIM(denoised_hsi_trans, test_label_hsi_trans) sam = SAM(denoised_hsi_trans, test_label_hsi_trans) #计算pnsr和ssim print("=====averPSNR:{:.4f}=====averSSIM:{:.4f}=====averSAM:{:.4f}". format(mpsnr, mssim, sam)) tb_writer.add_scalars("validation metrics", { 'average PSNR': mpsnr, 'average SSIM': mssim, 'avarage SAM': sam }, epoch) #通过这个我就可以看到,那个epoch的性能是最好的 #保存best模型 if mpsnr > best_psnr: best_psnr = mpsnr best_epoch = epoch best_iter = cur_step torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"{save_model_path}/hsid_rdn_4rdb_conise_l1_loss_600epoch_patchsize32_best.pth" ) print( "[epoch %d it %d PSNR: %.4f --- best_epoch %d best_iter %d Best_PSNR %.4f]" % (epoch, cur_step, psnr, best_epoch, best_iter, best_psnr)) print( "------------------------------------------------------------------" ) print("Epoch: {}\tTime: {:.4f}\tLoss: {:.4f}\tLearningRate {:.6f}". format(epoch, time.time() - epoch_start_time, gen_epoch_loss, INIT_LEARNING_RATE)) print( "------------------------------------------------------------------" ) #保存当前模型 torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict() }, os.path.join(save_model_path, "model_latest.pth")) tb_writer.close()
def train_model_multistage_lowlight(): device = DEVICE #准备数据 train_set = HsiCubicTrainDataset('./data/train_lowlight_patchsize32/') #print('trainset32 training example:', len(train_set32)) #train_set_64 = HsiCubicTrainDataset('./data/train_lowlight_patchsize64/') #train_set_list = [train_set32, train_set_64] #train_set = ConcatDataset(train_set_list) #里面的样本大小必须是一致的,否则会连接失败 print('total training example:', len(train_set)) train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=True) #加载测试label数据 mat_src_path = './data/test_lowlight/origin/soup_bigcorn_orange_1ms.mat' test_label_hsi = scio.loadmat(mat_src_path)['label'] #加载测试数据 batch_size = 1 test_data_dir = './data/test_lowlight/cubic/' test_set = HsiCubicLowlightTestDataset(test_data_dir) test_dataloader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=False) batch_size, channel, width, height = next(iter(test_dataloader))[0].shape band_num = len(test_dataloader) denoised_hsi = np.zeros((width, height, band_num)) #创建模型 net = MultiStageHSID(K) init_params(net) #net = nn.DataParallel(net).to(device) net = net.to(device) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE) scheduler = MultiStepLR(hsid_optimizer, milestones=[40, 60, 80], gamma=0.1) #定义loss 函数 #criterion = nn.MSELoss() global tb_writer tb_writer = get_summary_writer(log_dir='logs') gen_epoch_loss_list = [] cur_step = 0 first_batch = next(iter(train_loader)) best_psnr = 0 best_epoch = 0 best_iter = 0 start_epoch = 1 num_epoch = 100 for epoch in range(start_epoch, num_epoch + 1): epoch_start_time = time.time() scheduler.step() print(epoch, 'lr={:.6f}'.format(scheduler.get_last_lr()[0])) print(scheduler.get_lr()) gen_epoch_loss = 0 net.train() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for batch_idx, (noisy, cubic, label) in enumerate(train_loader): #print('batch_idx=', batch_idx) noisy = noisy.to(device) label = label.to(device) cubic = cubic.to(device) hsid_optimizer.zero_grad() #denoised_img = net(noisy, cubic) #loss = loss_fuction(denoised_img, label) residual = net(noisy, cubic) #loss = loss_fuction(residual, label-noisy) loss = np.sum([ loss_fuction(residual[j], label) for j in range(len(residual)) ]) loss.backward() # calcu gradient hsid_optimizer.step() # update parameter gen_epoch_loss += loss.item() if cur_step % display_step == 0: if cur_step > 0: print( f"Epoch {epoch}: Step {cur_step}: Batch_idx {batch_idx}: MSE loss: {loss.item()}" ) else: print("Pretrained initial state") tb_writer.add_scalar("MSE loss", loss.item(), cur_step) #step ++,每一次循环,每一个batch的处理,叫做一个step cur_step += 1 gen_epoch_loss_list.append(gen_epoch_loss) tb_writer.add_scalar("mse epoch loss", gen_epoch_loss, epoch) #scheduler.step() #print("Decaying learning rate to %g" % scheduler.get_last_lr()[0]) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"checkpoints/hsid_multistage_patchsize64_{epoch}.pth") #测试代码 net.eval() for batch_idx, (noisy_test, cubic_test, label_test) in enumerate(test_dataloader): noisy_test = noisy_test.type(torch.FloatTensor) label_test = label_test.type(torch.FloatTensor) cubic_test = cubic_test.type(torch.FloatTensor) noisy_test = noisy_test.to(DEVICE) label_test = label_test.to(DEVICE) cubic_test = cubic_test.to(DEVICE) with torch.no_grad(): residual = net(noisy_test, cubic_test) denoised_band = noisy_test + residual[0] denoised_band_numpy = denoised_band.cpu().numpy().astype( np.float32) denoised_band_numpy = np.squeeze(denoised_band_numpy) denoised_hsi[:, :, batch_idx] = denoised_band_numpy if batch_idx == 49: residual_squeezed = torch.squeeze(residual[0], axis=0) denoised_band_squeezed = torch.squeeze(denoised_band, axis=0) label_test_squeezed = torch.squeeze(label_test, axis=0) noisy_test_squeezed = torch.squeeze(noisy_test, axis=0) tb_writer.add_image(f"images/{epoch}_restored", denoised_band_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_residual", residual_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_label", label_test_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_noisy", noisy_test_squeezed, 1, dataformats='CHW') psnr = PSNR(denoised_hsi, test_label_hsi) ssim = SSIM(denoised_hsi, test_label_hsi) sam = SAM(denoised_hsi, test_label_hsi) #计算pnsr和ssim print("=====averPSNR:{:.3f}=====averSSIM:{:.4f}=====averSAM:{:.3f}". format(psnr, ssim, sam)) tb_writer.add_scalars("validation metrics", { 'average PSNR': psnr, 'average SSIM': ssim, 'avarage SAM': sam }, epoch) #通过这个我就可以看到,那个epoch的性能是最好的 #保存best模型 if psnr > best_psnr: best_psnr = psnr best_epoch = epoch best_iter = cur_step torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"checkpoints/hsid_multistage_patchsize64_best.pth") print( "[epoch %d it %d PSNR: %.4f --- best_epoch %d best_iter %d Best_PSNR %.4f]" % (epoch, cur_step, psnr, best_epoch, best_iter, best_psnr)) print( "------------------------------------------------------------------" ) print("Epoch: {}\tTime: {:.4f}\tLoss: {:.4f}\tLearningRate {:.6f}". format(epoch, time.time() - epoch_start_time, gen_epoch_loss, scheduler.get_lr()[0])) print( "------------------------------------------------------------------" ) #保存当前模型 torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict() }, os.path.join('./checkpoints', "model_latest.pth")) tb_writer.close()
def train_model_residual_lowlight_twostage_gan_best(): #设置超参数 batchsize = 128 init_lr = 0.001 K_adjacent_band = 36 display_step = 20 display_band = 20 is_resume = False lambda_recon = 10 start_epoch = 1 device = DEVICE #准备数据 train_set = HsiCubicTrainDataset('./data/train_lowlight/') print('total training example:', len(train_set)) train_loader = DataLoader(dataset=train_set, batch_size=batchsize, shuffle=True) #加载测试label数据 mat_src_path = './data/test_lowlight/origin/soup_bigcorn_orange_1ms.mat' test_label_hsi = scio.loadmat(mat_src_path)['label'] #加载测试数据 test_batch_size = 1 test_data_dir = './data/test_lowlight/cubic/' test_set = HsiCubicLowlightTestDataset(test_data_dir) test_dataloader = DataLoader(dataset=test_set, batch_size=test_batch_size, shuffle=False) batch_size, channel, width, height = next(iter(test_dataloader))[0].shape band_num = len(test_dataloader) denoised_hsi = np.zeros((width, height, band_num)) #创建模型 net = HSIDDenseNetTwoStage(K_adjacent_band) init_params(net) #net = nn.DataParallel(net).to(device) net = net.to(device) #创建discriminator disc = DiscriminatorABC(2, 4) init_params(disc) disc = disc.to(device) disc_opt = torch.optim.Adam(disc.parameters(), lr=init_lr) num_epoch = 100 print('epoch count == ', num_epoch) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam(net.parameters(), lr=init_lr) #Scheduler scheduler = MultiStepLR(hsid_optimizer, milestones=[40, 60, 80], gamma=0.1) warmup_epochs = 3 #scheduler_cosine = optim.lr_scheduler.CosineAnnealingLR(hsid_optimizer, num_epoch-warmup_epochs+40, eta_min=1e-7) #scheduler = GradualWarmupScheduler(hsid_optimizer, multiplier=1, total_epoch=warmup_epochs, after_scheduler=scheduler_cosine) #scheduler.step() #唤醒训练 if is_resume: model_dir = './checkpoints' path_chk_rest = dir_utils.get_last_path(model_dir, 'model_latest.pth') model_utils.load_checkpoint(net, path_chk_rest) start_epoch = model_utils.load_start_epoch(path_chk_rest) + 1 model_utils.load_optim(hsid_optimizer, path_chk_rest) model_utils.load_disc_checkpoint(disc, path_chk_rest) model_utils.load_disc_optim(disc_opt, path_chk_rest) for i in range(1, start_epoch): scheduler.step() new_lr = scheduler.get_lr()[0] print( '------------------------------------------------------------------------------' ) print("==> Resuming Training with learning rate:", new_lr) print( '------------------------------------------------------------------------------' ) #定义loss 函数 #criterion = nn.MSELoss() global tb_writer tb_writer = get_summary_writer(log_dir='logs') gen_epoch_loss_list = [] cur_step = 0 first_batch = next(iter(train_loader)) best_psnr = 0 best_epoch = 0 best_iter = 0 for epoch in range(start_epoch, num_epoch + 1): epoch_start_time = time.time() scheduler.step() #print(epoch, 'lr={:.6f}'.format(scheduler.get_last_lr()[0])) print('epoch = ', epoch, 'lr={:.6f}'.format(scheduler.get_lr()[0])) print(scheduler.get_lr()) gen_epoch_loss = 0 net.train() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for batch_idx, (noisy, cubic, label) in enumerate(train_loader): #print('batch_idx=', batch_idx) noisy = noisy.to(device) label = label.to(device) cubic = cubic.to(device) ### Update discriminator ### disc_opt.zero_grad( ) # Zero out the gradient before backpropagation with torch.no_grad(): fake, fake_stage2 = net(noisy, cubic) #print('noisy shape =', noisy.shape, fake_stage2.shape) #fake.detach() disc_fake_hat = disc(fake_stage2.detach() + noisy, noisy) # Detach generator disc_fake_loss = adv_criterion(disc_fake_hat, torch.zeros_like(disc_fake_hat)) disc_real_hat = disc(label, noisy) disc_real_loss = adv_criterion(disc_real_hat, torch.ones_like(disc_real_hat)) disc_loss = (disc_fake_loss + disc_real_loss) / 2 disc_loss.backward(retain_graph=True) # Update gradients disc_opt.step() # Update optimizer ### Update generator ### hsid_optimizer.zero_grad() #denoised_img = net(noisy, cubic) #loss = loss_fuction(denoised_img, label) residual, residual_stage2 = net(noisy, cubic) disc_fake_hat = disc(residual_stage2 + noisy, noisy) gen_adv_loss = adv_criterion(disc_fake_hat, torch.ones_like(disc_fake_hat)) alpha = 0.2 beta = 0.2 rec_loss = beta * (alpha*loss_fuction(residual, label-noisy) + (1-alpha) * recon_criterion(residual, label-noisy)) \ + (1-beta) * (alpha*loss_fuction(residual_stage2, label-noisy) + (1-alpha) * recon_criterion(residual_stage2, label-noisy)) loss = gen_adv_loss + lambda_recon * rec_loss loss.backward() # calcu gradient hsid_optimizer.step() # update parameter gen_epoch_loss += loss.item() if cur_step % display_step == 0: if cur_step > 0: print( f"Epoch {epoch}: Step {cur_step}: Batch_idx {batch_idx}: MSE loss: {loss.item()}" ) print( f"rec_loss = {rec_loss.item()}, gen_adv_loss = {gen_adv_loss.item()}" ) else: print("Pretrained initial state") tb_writer.add_scalar("MSE loss", loss.item(), cur_step) #step ++,每一次循环,每一个batch的处理,叫做一个step cur_step += 1 gen_epoch_loss_list.append(gen_epoch_loss) tb_writer.add_scalar("mse epoch loss", gen_epoch_loss, epoch) #scheduler.step() #print("Decaying learning rate to %g" % scheduler.get_last_lr()[0]) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), 'disc': disc.state_dict(), 'disc_opt': disc_opt.state_dict() }, f"checkpoints/two_stage_hsid_dense_gan_{epoch}.pth") #测试代码 net.eval() for batch_idx, (noisy_test, cubic_test, label_test) in enumerate(test_dataloader): noisy_test = noisy_test.type(torch.FloatTensor) label_test = label_test.type(torch.FloatTensor) cubic_test = cubic_test.type(torch.FloatTensor) noisy_test = noisy_test.to(DEVICE) label_test = label_test.to(DEVICE) cubic_test = cubic_test.to(DEVICE) with torch.no_grad(): residual, residual_stage2 = net(noisy_test, cubic_test) denoised_band = noisy_test + residual_stage2 denoised_band_numpy = denoised_band.cpu().numpy().astype( np.float32) denoised_band_numpy = np.squeeze(denoised_band_numpy) denoised_hsi[:, :, batch_idx] = denoised_band_numpy if batch_idx == 49: residual_squeezed = torch.squeeze(residual, axis=0) residual_stage2_squeezed = torch.squeeze(residual_stage2, axis=0) denoised_band_squeezed = torch.squeeze(denoised_band, axis=0) label_test_squeezed = torch.squeeze(label_test, axis=0) noisy_test_squeezed = torch.squeeze(noisy_test, axis=0) tb_writer.add_image(f"images/{epoch}_restored", denoised_band_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_residual", residual_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_residual_stage2", residual_stage2_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_label", label_test_squeezed, 1, dataformats='CHW') tb_writer.add_image(f"images/{epoch}_noisy", noisy_test_squeezed, 1, dataformats='CHW') psnr = PSNR(denoised_hsi, test_label_hsi) ssim = SSIM(denoised_hsi, test_label_hsi) sam = SAM(denoised_hsi, test_label_hsi) #计算pnsr和ssim print("=====averPSNR:{:.3f}=====averSSIM:{:.4f}=====averSAM:{:.3f}". format(psnr, ssim, sam)) tb_writer.add_scalars("validation metrics", { 'average PSNR': psnr, 'average SSIM': ssim, 'avarage SAM': sam }, epoch) #通过这个我就可以看到,那个epoch的性能是最好的 #保存best模型 if psnr > best_psnr: best_psnr = psnr best_epoch = epoch best_iter = cur_step torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), 'disc': disc.state_dict(), 'disc_opt': disc_opt.state_dict() }, f"checkpoints/two_stage_hsid_dense_gan_best.pth") print( "[epoch %d it %d PSNR: %.4f --- best_epoch %d best_iter %d Best_PSNR %.4f]" % (epoch, cur_step, psnr, best_epoch, best_iter, best_psnr)) print( "------------------------------------------------------------------" ) print("Epoch: {}\tTime: {:.4f}\tLoss: {:.4f}\tLearningRate {:.6f}". format(epoch, time.time() - epoch_start_time, gen_epoch_loss, scheduler.get_lr()[0])) print( "------------------------------------------------------------------" ) torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), 'disc': disc.state_dict(), 'disc_opt': disc_opt.state_dict() }, os.path.join('./checkpoints', "model_latest.pth")) tb_writer.close()
def train_model_residual_lowlight_rdn(): device = DEVICE #准备数据 train = np.load('./data/denoise/train_washington8.npy') train = train.transpose((2, 1, 0)) test = np.load('./data/denoise/train_washington8.npy') #test=test.transpose((2,1,0)) test = test.transpose((2, 1, 0)) #将通道维放在最前面 save_model_path = './checkpoints/hsirnd_denoise_l1loss' if not os.path.exists(save_model_path): os.mkdir(save_model_path) #创建模型 net = HSIRDNECA_Denoise(K) init_params(net) net = nn.DataParallel(net).to(device) #net = net.to(device) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE) scheduler = MultiStepLR(hsid_optimizer, milestones=[200, 400], gamma=0.5) #定义loss 函数 #criterion = nn.MSELoss() gen_epoch_loss_list = [] cur_step = 0 best_psnr = 0 best_epoch = 0 best_iter = 0 start_epoch = 1 num_epoch = 600 mpsnr_list = [] for epoch in range(start_epoch, num_epoch + 1): epoch_start_time = time.time() scheduler.step() print('epoch = ', epoch, 'lr={:.6f}'.format(scheduler.get_lr()[0])) print(scheduler.get_lr()) gen_epoch_loss = 0 net.train() channels = 191 # 191 channels data_patches, data_cubic_patches = datagenerator(train, channels) data_patches = torch.from_numpy(data_patches.transpose(( 0, 3, 1, 2, ))) data_cubic_patches = torch.from_numpy( data_cubic_patches.transpose((0, 4, 1, 2, 3))) DDataset = DenoisingDataset(data_patches, data_cubic_patches, SIGMA) print('yes') DLoader = DataLoader(dataset=DDataset, batch_size=BATCH_SIZE, shuffle=True) # loader出问题了 epoch_loss = 0 start_time = time.time() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for step, x_y in enumerate(DLoader): #print('batch_idx=', batch_idx) batch_x_noise, batch_y_noise, batch_x = x_y[0], x_y[1], x_y[2] batch_x_noise = batch_x_noise.to(device) batch_y_noise = batch_y_noise.to(device) batch_x = batch_x.to(device) hsid_optimizer.zero_grad() #denoised_img = net(noisy, cubic) #loss = loss_fuction(denoised_img, label) residual = net(batch_x_noise, batch_y_noise) alpha = 0.8 loss = recon_criterion(residual, batch_x - batch_x_noise) #loss = alpha*recon_criterion(residual, label-noisy) + (1-alpha)*loss_function_mse(residual, label-noisy) #loss = recon_criterion(residual, label-noisy) loss.backward() # calcu gradient hsid_optimizer.step() # update parameter if step % 10 == 0: print('%4d %4d / %4d loss = %2.8f' % (epoch + 1, step, data_patches.size(0) // BATCH_SIZE, loss.item() / BATCH_SIZE)) #scheduler.step() #print("Decaying learning rate to %g" % scheduler.get_last_lr()[0]) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"{save_model_path}/hsid_rdn_eca_l1_loss_600epoch_patchsize32_{epoch}.pth" ) #测试代码 net.eval() """ channel_s = 191 # 设置多少波段 data_patches, data_cubic_patches = datagenerator(test, channel_s) data_patches = torch.from_numpy(data_patches.transpose((0, 3, 1, 2,))) data_cubic_patches = torch.from_numpy(data_cubic_patches.transpose((0, 4, 1, 2, 3))) DDataset = DenoisingDataset(data_patches, data_cubic_patches, SIGMA) DLoader = DataLoader(dataset=DDataset, batch_size=BATCH_SIZE, shuffle=True) epoch_loss = 0 for step, x_y in enumerate(DLoader): batch_x_noise, batch_y_noise, batch_x = x_y[0], x_y[1], x_y[2] batch_x_noise = batch_x_noise.to(DEVICE) batch_y_noise = batch_y_noise.to(DEVICE) batch_x = batch_x.to(DEVICE) residual = net(batch_x_noise, batch_y_noise) loss = loss_fuction(residual, batch_x-batch_x_noise) epoch_loss += loss.item() if step % 10 == 0: print('%4d %4d / %4d test loss = %2.4f' % ( epoch + 1, step, data_patches.size(0) // BATCH_SIZE, loss.item() / BATCH_SIZE)) """ #加载数据 test_data_dir = './data/denoise/test/' test_set = HsiTrainDataset(test_data_dir) test_dataloader = DataLoader(test_set, batch_size=1, shuffle=False) #指定结果输出路径 test_result_output_path = './data/denoise/testresult/' if not os.path.exists(test_result_output_path): os.makedirs(test_result_output_path) #逐个通道的去噪 """ 分配一个numpy数组,存储去噪后的结果 遍历所有通道, 对于每个通道,通过get_adjacent_spectral_bands获取其相邻的K个通道 调用hsid进行预测 将预测到的residual和输入的noise加起来,得到输出band 将去噪后的结果保存成mat结构 """ psnr_list = [] for batch_idx, (noisy, label) in enumerate(test_dataloader): noisy = noisy.type(torch.FloatTensor) label = label.type(torch.FloatTensor) batch_size, width, height, band_num = noisy.shape denoised_hsi = np.zeros((width, height, band_num)) noisy = noisy.to(DEVICE) label = label.to(DEVICE) with torch.no_grad(): for i in range(band_num): #遍历每个band去处理 current_noisy_band = noisy[:, :, :, i] current_noisy_band = current_noisy_band[:, None] adj_spectral_bands = get_adjacent_spectral_bands( noisy, K, i) #adj_spectral_bands = torch.transpose(adj_spectral_bands,3,1) #将通道数置换到第二维 adj_spectral_bands = adj_spectral_bands.permute(0, 3, 1, 2) adj_spectral_bands_unsqueezed = adj_spectral_bands.unsqueeze( 1) #print(current_noisy_band.shape, adj_spectral_bands.shape) residual = net(current_noisy_band, adj_spectral_bands_unsqueezed) denoised_band = residual + current_noisy_band denoised_band_numpy = denoised_band.cpu().numpy().astype( np.float32) denoised_band_numpy = np.squeeze(denoised_band_numpy) denoised_hsi[:, :, i] += denoised_band_numpy test_label_current_band = label[:, :, :, i] label_band_numpy = test_label_current_band.cpu().numpy( ).astype(np.float32) label_band_numpy = np.squeeze(label_band_numpy) #print(denoised_band_numpy.shape, label_band_numpy.shape, label.shape) psnr = PSNR(denoised_band_numpy, label_band_numpy) psnr_list.append(psnr) mpsnr = np.mean(psnr_list) mpsnr_list.append(mpsnr) denoised_hsi_trans = denoised_hsi.transpose(2, 0, 1) test_label_hsi_trans = np.squeeze(label.cpu().numpy().astype( np.float32)).transpose(2, 0, 1) mssim = SSIM(denoised_hsi_trans, test_label_hsi_trans) sam = SAM(denoised_hsi_trans, test_label_hsi_trans) #计算pnsr和ssim print( "=====averPSNR:{:.3f}=====averSSIM:{:.4f}=====averSAM:{:.3f}". format(mpsnr, mssim, sam)) #保存best模型 if mpsnr > best_psnr: best_psnr = mpsnr best_epoch = epoch best_iter = cur_step torch.save( { 'epoch': epoch, 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"{save_model_path}/hsid_rdn_eca_l1_loss_600epoch_patchsize32_best.pth" ) print( "[epoch %d it %d PSNR: %.4f --- best_epoch %d best_iter %d Best_PSNR %.4f]" % (epoch, cur_step, mpsnr, best_epoch, best_iter, best_psnr)) print( "------------------------------------------------------------------" ) print("Epoch: {}\tTime: {:.4f}\tLoss: {:.4f}\tLearningRate {:.6f}". format(epoch, time.time() - epoch_start_time, gen_epoch_loss, INIT_LEARNING_RATE)) print( "------------------------------------------------------------------" )
def main(): device = DEVICE #准备数据 train_set = HsiTrainDataset('./data/train/') train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=True) #创建模型 net = HSID_1x3(K) init_params(net) net = nn.DataParallel(net).to(device) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam( net.parameters(), lr=INIT_LEARNING_RATE) #betas default value 就是0.9和0.999 scheduler = MultiStepLR(hsid_optimizer, milestones=[15, 30, 45], gamma=0.25) #定义loss 函数 #criterion = nn.MSELoss() global tb_writer tb_writer = get_summary_writer(log_dir='logs') gen_minibatch_loss_list = [] gen_epoch_loss_list = [] cur_step = 0 first_batch = next(iter(train_loader)) for epoch in range(NUM_EPOCHS): gen_epoch_loss = 0 net.train() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for batch_idx, (noisy, label) in enumerate(train_loader): noisy = noisy.to(device) label = label.to(device) batch_size, height, width, band_num = noisy.shape """" our method traverses all the bands through one-by-one mode, which simultaneously employing spatial–spectral information with spatial and spatial–spectral filters, respectively """ band_loss = 0 for i in range(band_num): #遍历每个band去处理 single_noisy_band = noisy[:, :, :, i] single_noisy_band_cloned = single_noisy_band[:, None].clone() single_label_band = label[:, :, :, i] single_label_band_cloned = single_label_band[:, None].clone() adj_spectral_bands = get_adjacent_spectral_bands(noisy, K, i) #print('adj_spectral_bands.shape =', adj_spectral_bands.shape) #print(type(adj_spectral_bands)) adj_spectral_bands_transposed = torch.transpose( adj_spectral_bands, 3, 1).clone() #print('transposed adj_spectral_bands.shape =', adj_spectral_bands.shape) #print(type(adj_spectral_bands)) denoised_img = net(single_noisy_band_cloned, adj_spectral_bands_transposed) loss = loss_fuction(single_label_band_cloned, denoised_img) hsid_optimizer.zero_grad() loss.backward() # calcu gradient hsid_optimizer.step() # update parameter ## Logging band_loss += loss.item() if i % 20 == 0: print( f"Epoch {epoch}: Step {cur_step}: bandnum {i}: band MSE loss: {loss.item()}" ) gen_minibatch_loss_list.append(band_loss) gen_epoch_loss += band_loss if cur_step % display_step == 0: if cur_step > 0: print( f"Epoch {epoch}: Step {cur_step}: MSE loss: {band_loss}" ) else: print("Pretrained initial state") tb_writer.add_scalar("MSE loss", band_loss, cur_step) #step ++,每一次循环,每一个batch的处理,叫做一个step cur_step += 1 scheduler.step() print("Decaying learning rate to %g" % scheduler.get_lr()[0]) gen_epoch_loss_list.append(gen_epoch_loss) tb_writer.add_scalar("mse epoch loss", gen_epoch_loss, epoch) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"checkpoints/hsid_1x3{epoch}.pth") tb_writer.close()
def train_model_residual_lowlight_twostage_unet(): learning_rate = INIT_LEARNING_RATE * 0.5 device = DEVICE #准备数据 train_set = HsiCubicTrainDataset('./data/train_lowlight/') print('total training example:', len(train_set)) train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=True) #加载测试label数据 mat_src_path = './data/test_lowlight/origin/soup_bigcorn_orange_1ms.mat' test_label_hsi = scio.loadmat(mat_src_path)['label'] #加载测试数据 batch_size = 1 test_data_dir = './data/test_lowlight/cubic/' test_set = HsiCubicLowlightTestDataset(test_data_dir) test_dataloader = DataLoader(dataset=test_set, batch_size=batch_size, shuffle=False) batch_size, channel, width, height = next(iter(test_dataloader))[0].shape band_num = len(test_dataloader) denoised_hsi = np.zeros((width, height, band_num)) #创建模型 net = TwoStageHSIDWithUNet(K) init_params(net) #net = nn.DataParallel(net).to(device) net = net.to(device) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam(net.parameters(), lr=learning_rate) scheduler = MultiStepLR(hsid_optimizer, milestones=[40, 60, 80], gamma=0.1) #定义loss 函数 #criterion = nn.MSELoss() global tb_writer tb_writer = get_summary_writer(log_dir='logs') gen_epoch_loss_list = [] cur_step = 0 first_batch = next(iter(train_loader)) for epoch in range(NUM_EPOCHS): scheduler.step() print(epoch, 'lr={:.6f}'.format(scheduler.get_last_lr()[0])) gen_epoch_loss = 0 net.train() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for batch_idx, (noisy, cubic, label) in enumerate(train_loader): #print('batch_idx=', batch_idx) noisy = noisy.to(device) label = label.to(device) cubic = cubic.to(device) hsid_optimizer.zero_grad() #denoised_img = net(noisy, cubic) #loss = loss_fuction(denoised_img, label) residual, residual_stage2 = net(noisy, cubic) loss = loss_function_with_tvloss( residual, label - noisy) + loss_function_with_tvloss( residual_stage2, label - noisy) loss.backward() # calcu gradient hsid_optimizer.step() # update parameter gen_epoch_loss += loss.item() if cur_step % display_step == 0: if cur_step > 0: print( f"Epoch {epoch}: Step {cur_step}: Batch_idx {batch_idx}: MSE loss: {loss.item()}" ) else: print("Pretrained initial state") tb_writer.add_scalar("MSE loss", loss.item(), cur_step) #step ++,每一次循环,每一个batch的处理,叫做一个step cur_step += 1 gen_epoch_loss_list.append(gen_epoch_loss) tb_writer.add_scalar("mse epoch loss", gen_epoch_loss, epoch) #scheduler.step() #print("Decaying learning rate to %g" % scheduler.get_last_lr()[0]) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"checkpoints/two_stage_unet_hsid_{epoch}.pth") #测试代码 net.eval() for batch_idx, (noisy_test, cubic_test, label_test) in enumerate(test_dataloader): noisy_test = noisy_test.type(torch.FloatTensor) label_test = label_test.type(torch.FloatTensor) cubic_test = cubic_test.type(torch.FloatTensor) noisy_test = noisy_test.to(DEVICE) label_test = label_test.to(DEVICE) cubic_test = cubic_test.to(DEVICE) with torch.no_grad(): residual, residual_stage2 = net(noisy_test, cubic_test) denoised_band = noisy_test + residual_stage2 denoised_band_numpy = denoised_band.cpu().numpy().astype( np.float32) denoised_band_numpy = np.squeeze(denoised_band_numpy) denoised_hsi[:, :, batch_idx] = denoised_band_numpy psnr = PSNR(denoised_hsi, test_label_hsi) ssim = SSIM(denoised_hsi, test_label_hsi) sam = SAM(denoised_hsi, test_label_hsi) #计算pnsr和ssim print("=====averPSNR:{:.3f}=====averSSIM:{:.4f}=====averSAM:{:.3f}". format(psnr, ssim, sam)) tb_writer.add_scalars("validation metrics", { 'average PSNR': psnr, 'average SSIM': ssim, 'avarage SAM': sam }, epoch) #通过这个我就可以看到,那个epoch的性能是最好的 tb_writer.close()
def train_model(): device = DEVICE #准备数据 train_set = HsiCubicTrainDataset('./data/train_cubic/') train_loader = DataLoader(dataset=train_set, batch_size=BATCH_SIZE, shuffle=True) #加载测试label数据 test_label_hsi = np.load('./data/origin/test_washington.npy') #加载测试数据 test_data_dir = './data/test_level25/' test_set = HsiTrainDataset(test_data_dir) test_dataloader = DataLoader(test_set, batch_size=1, shuffle=False) #创建模型 net = HSID_1x3(K) init_params(net) net = nn.DataParallel(net).to(device) #创建优化器 #hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE, betas=(0.9, 0,999)) hsid_optimizer = optim.Adam(net.parameters(), lr=INIT_LEARNING_RATE) scheduler = MultiStepLR(hsid_optimizer, milestones=[15, 30, 45], gamma=0.25) #定义loss 函数 #criterion = nn.MSELoss() global tb_writer tb_writer = get_summary_writer(log_dir='logs') gen_epoch_loss_list = [] cur_step = 0 first_batch = next(iter(train_loader)) for epoch in range(NUM_EPOCHS): gen_epoch_loss = 0 net.train() #for batch_idx, (noisy, label) in enumerate([first_batch] * 300): for batch_idx, (noisy, cubic, label) in enumerate(train_loader): noisy = noisy.to(device) label = label.to(device) cubic = cubic.to(device) hsid_optimizer.zero_grad() denoised_img = net(noisy, cubic) loss = loss_fuction(denoised_img, label) loss.backward() # calcu gradient hsid_optimizer.step() # update parameter gen_epoch_loss += loss.item() if cur_step % display_step == 0: if cur_step > 0: print( f"Epoch {epoch}: Step {cur_step}: MSE loss: {loss.item()}" ) else: print("Pretrained initial state") tb_writer.add_scalar("MSE loss", loss.item(), cur_step) #step ++,每一次循环,每一个batch的处理,叫做一个step cur_step += 1 gen_epoch_loss_list.append(gen_epoch_loss) tb_writer.add_scalar("mse epoch loss", gen_epoch_loss, epoch) scheduler.step() print("Decaying learning rate to %g" % scheduler.get_last_lr()[0]) torch.save( { 'gen': net.state_dict(), 'gen_opt': hsid_optimizer.state_dict(), }, f"checkpoints/hsid_{epoch}.pth") #预测代码 net.eval() for batch_idx, (noisy, label) in enumerate(test_dataloader): noisy = noisy.type(torch.FloatTensor) label = label.type(torch.FloatTensor) batch_size, width, height, band_num = noisy.shape denoised_hsi = np.zeros((width, height, band_num)) noisy = noisy.to(DEVICE) label = label.to(DEVICE) with torch.no_grad(): for i in range(band_num): #遍历每个band去处理 current_noisy_band = noisy[:, :, :, i] current_noisy_band = current_noisy_band[:, None] adj_spectral_bands = get_adjacent_spectral_bands( noisy, K, i) # shape: batch_size, width, height, band_num adj_spectral_bands = torch.transpose( adj_spectral_bands, 3, 1 ) #交换第一维和第三维 ,shape: batch_size, band_num, height, width denoised_band = net(current_noisy_band, adj_spectral_bands) denoised_band_numpy = denoised_band.cpu().numpy().astype( np.float32) denoised_band_numpy = np.squeeze(denoised_band_numpy) denoised_hsi[:, :, i] = denoised_band_numpy psnr = PSNR(denoised_hsi, test_label_hsi) ssim = SSIM(denoised_hsi, test_label_hsi) sam = SAM(denoised_hsi, test_label_hsi) #计算pnsr和ssim print("=====averPSNR:{:.3f}=====averSSIM:{:.4f}=====averSAM:{:.3f}". format(psnr, ssim, sam)) tb_writer.close()