def initialize(self, opt): BaseModel.initialize(self, opt) self.isTrain = opt.isTrain # specify the training losses you want to print out. The program will call base_model.get_current_losses if self.isTrain: if self.train_phase == 'generator': self.model_names = ['G'] self.loss_names = ['G_I_L1','G_I_L2','SSIM', 'PSNR', 'vgg'] else: self.model_names = ['G', 'D'] self.loss_names = ['G_GAN', 'G_I_L1', 'G_I_L2', 'D_GAN_fake', 'D_GAN_real', 'SSIM', 'PSNR', 'vgg'] else: # during test time, only load Gs self.model_names = ['G'] self.loss_names = ['SSIM', 'PSNR'] # specify the images you want to save/display. The program will call base_model.get_current_visuals self.visual_names = ['real_A', 'fake_B', 'real_B'] self.netG = networks.define_G(self.opt, opt.input_nc, opt.output_nc, opt.ngf, opt.which_model_netG, opt.norm, not opt.no_dropout, opt.init_type, opt.init_gain, self.gpu_ids) self.criterionL1 = torch.nn.L1Loss() self.criterionMSE = torch.nn.MSELoss() self.ssim_loss = pytorch_msssim.SSIM(val_range=1) if opt.use_vgg: self.perceptual = losses.PerceptualLoss() self.perceptual.initialize(self.criterionMSE) if self.isTrain: self.optimizers = [] self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_G) if self.isTrain and self.train_phase == 'together': self.no_wgan = opt.no_wgan self.no_wgan_gp = opt.no_wgan_gp if self.no_wgan_gp == False: self.disc_step = opt.disc_step else: self.disc_step = 1 self.disc_model = opt.disc_model use_sigmoid = opt.no_lsgan if opt.disc_model=='pix2pix': self.netD = networks.define_D(opt,opt.input_nc + opt.output_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids) if opt.disc_model=='traditional': self.netD = networks.define_D(opt,opt.output_nc, opt.ndf, opt.which_model_netD, opt.n_layers_D, opt.norm, use_sigmoid, opt.init_type, opt.init_gain, self.gpu_ids) self.loss_wgan_gp = opt.loss_wgan_gp self.fake_pool = ImagePool(opt.pool_size) self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, use_l1=not opt.no_l1gan).to(self.device) self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999)) self.optimizers.append(self.optimizer_D)
def compute_loss(predicts, targets, weights, w_idx=0): # targets = Variable(targets.data.clone(), requires_grad=False) mse = nn.MSELoss() ssim = pytorch_msssim.SSIM() loss_mse = mse(predicts, targets) loss_ssim = 1 - ssim(predicts, targets) loss = loss_mse + weights[w_idx] * loss_ssim return loss
def __init__(self, args): super(Loss, self).__init__() print('Preparing loss function:') self.loss = [] self.loss_module = nn.ModuleList() for loss in args.loss.split('+'): weight, loss_type = loss.split('*') if loss_type == 'MSE': loss_function = nn.MSELoss() elif loss_type == 'Huber': loss_function = HuberLoss(delta=.5) elif loss_type == 'L1': loss_function = nn.L1Loss() elif loss_type.find('VGG') >= 0: loss_function = VGG(loss_type[3:]) elif loss_type == 'SSIM': loss_function = pytorch_msssim.SSIM(val_range=1.) elif loss_type.find('GAN') >= 0: loss_function = Adversarial(args, loss_type) self.loss.append({ 'type': loss_type, 'weight': float(weight), 'function': loss_function }) if loss_type.find('GAN') >= 0 >= 0: self.loss.append({ 'type': 'DIS', 'weight': 1, 'function': None }) if len(self.loss) > 1: self.loss.append({'type': 'Total', 'weight': 0, 'function': None}) for l in self.loss: if l['function'] is not None: print('{:.3f} * {}'.format(l['weight'], l['type'])) self.loss_module.append(l['function']) device = torch.device('cuda' if args.cuda else 'cpu') self.loss_module.to(device) #if args.precision == 'half': self.loss_module.half() if args.cuda: # and args.n_GPUs > 1: self.loss_module = nn.DataParallel(self.loss_module)
# Create Model model = autoencoder.ConvolutionalAE( max_filters=max_filters, num_layers=num_layers, input_image_dimensions=image_size, latent_dim=latent_dim, small_conv=small_conv, ) model.to(device) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) ssim_module = None if use_ssim_loss: ssim_module = pytorch_msssim.SSIM(data_range=1.0, win_size=11, win_sigma=1.5, K=(0.01, 0.03)) ################################################################################ ################################### Training ################################### ################################################################################ # Train all_samples = [] all_train_loss = [] all_val_loss = [] # Get an initial "epoch 0" sample model.eval() with torch.no_grad(): epoch_sample = model(sample.to(device))
def train(args): print('Number of GPUs available: ' + str(torch.cuda.device_count())) model = nn.DataParallel(CAEP(num_resblocks).cuda()) print('Done Setup Model.') dataset = BSDS500Crop128(args.dataset_path) dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=args.shuffle, num_workers=args.num_workers) testset = Kodak(args.testset_path) testloader = DataLoader(testset, batch_size=testset.__len__(), num_workers=args.num_workers) print( f"Done Setup Training DataLoader: {len(dataloader)} batches of size {args.batch_size}" ) print(f"Done Setup Testing DataLoader: {len(testset)} Images") MSE = nn.MSELoss() SSIM = pytorch_msssim.SSIM().cuda() MSSSIM = pytorch_msssim.MSSSIM().cuda() optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=1e-10) scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=10, verbose=True, ) writer = SummaryWriter(log_dir=f'TBXLog/{args.exp_name}') # ADMM variables Z = torch.zeros(16, 32, 32).cuda() U = torch.zeros(16, 32, 32).cuda() Z.requires_grad = False U.requires_grad = False if args.load != '': pretrained_state_dict = torch.load(f"./chkpt/{args.load}/model.state") current_state_dict = model.state_dict() current_state_dict.update(pretrained_state_dict) model.load_state_dict(current_state_dict) # Z = torch.load(f"./chkpt/{args.load}/Z.state") # U = torch.load(f"./chkpt/{args.load}/U.state") if args.load == args.exp_name: optimizer.load_state_dict( torch.load(f"./chkpt/{args.load}/opt.state")) scheduler.load_state_dict( torch.load(f"./chkpt/{args.load}/lr.state")) print('Model Params Loaded.') model.train() for ei in range(args.res_epoch + 1, args.res_epoch + args.num_epochs + 1): # train train_loss = 0 train_ssim = 0 train_msssim = 0 train_psnr = 0 train_peanalty = 0 train_bpp = 0 avg_c = torch.zeros(16, 32, 32).cuda() avg_c.requires_grad = False for bi, crop in enumerate(dataloader): x = crop.cuda() y, c = model(x) psnr = compute_psnr(x, y) mse = MSE(y, x) ssim = SSIM(x, y) msssim = MSSSIM(x, y) mix = 1000 * (1 - msssim) + 1000 * (1 - ssim) + 1e4 * mse + (45 - psnr) # ADMM Step 1 peanalty = rho / 2 * torch.norm(c - Z + U, 2) bpp = compute_bpp(c, x.shape[0], 'crop', save=False) avg_c += torch.mean(c.detach() / (len(dataloader) * args.admm_every), dim=0) loss = mix + peanalty if ei == 1 and args.load != args.exp_name: loss = 1e5 * mse # warm up optimizer.zero_grad() loss.backward() optimizer.step() print( '[%3d/%3d][%5d/%5d] Loss: %f, SSIM: %f, MSSSIM: %f, PSNR: %f, Norm of Code: %f, BPP: %2f' % (ei, args.num_epochs + args.res_epoch, bi, len(dataloader), loss, ssim, msssim, psnr, peanalty, bpp)) writer.add_scalar('batch_train/loss', loss, ei * len(dataloader) + bi) writer.add_scalar('batch_train/ssim', ssim, ei * len(dataloader) + bi) writer.add_scalar('batch_train/msssim', msssim, ei * len(dataloader) + bi) writer.add_scalar('batch_train/psnr', psnr, ei * len(dataloader) + bi) writer.add_scalar('batch_train/norm', peanalty, ei * len(dataloader) + bi) writer.add_scalar('batch_train/bpp', bpp, ei * len(dataloader) + bi) train_loss += loss.item() / len(dataloader) train_ssim += ssim.item() / len(dataloader) train_msssim += msssim.item() / len(dataloader) train_psnr += psnr.item() / len(dataloader) train_peanalty += peanalty.item() / len(dataloader) train_bpp += bpp / len(dataloader) writer.add_scalar('epoch_train/loss', train_loss, ei) writer.add_scalar('epoch_train/ssim', train_ssim, ei) writer.add_scalar('epoch_train/msssim', train_msssim, ei) writer.add_scalar('epoch_train/psnr', train_psnr, ei) writer.add_scalar('epoch_train/norm', train_peanalty, ei) writer.add_scalar('epoch_train/bpp', train_bpp, ei) if ei % args.admm_every == args.admm_every - 1: # ADMM Step 2 Z = (avg_c + U).masked_fill_((torch.Tensor( np.argsort((avg_c + U).data.cpu().numpy(), axis=None)) >= int( (1 - pruning_ratio) * 16 * 32 * 32)).view(16, 32, 32).cuda(), value=0) # ADMM Step 3 U += avg_c - Z # test model.eval() val_loss = 0 val_ssim = 0 val_msssim = 0 val_psnr = 0 val_peanalty = 0 val_bpp = 0 for bi, (img, patches, _) in enumerate(testloader): avg_loss = 0 avg_ssim = 0 avg_msssim = 0 avg_psnr = 0 avg_peanalty = 0 avg_bpp = 0 for i in range(6): for j in range(4): x = torch.Tensor(patches[:, i, j, :, :, :]).cuda() y, c = model(x) psnr = compute_psnr(x, y) mse = MSE(y, x) ssim = SSIM(x, y) msssim = MSSSIM(x, y) mix = 1000 * (1 - msssim) + 1000 * ( 1 - ssim) + 1e4 * mse + (45 - psnr) peanalty = rho / 2 * torch.norm(c - Z + U, 2) bpp = compute_bpp(c, x.shape[0], f'Kodak_patches_{i}_{j}', save=True) loss = mix + peanalty avg_loss += loss.item() / 24 avg_ssim += ssim.item() / 24 avg_msssim += msssim.item() / 24 avg_psnr += psnr.item() / 24 avg_peanalty += peanalty.item() / 24 avg_bpp += bpp / 24 save_kodak_img(model, img, 0, patches, writer, ei) save_kodak_img(model, img, 10, patches, writer, ei) save_kodak_img(model, img, 20, patches, writer, ei) val_loss += avg_loss val_ssim += avg_ssim val_msssim += avg_msssim val_psnr += avg_psnr val_peanalty += avg_peanalty val_bpp += avg_bpp print( '*Kodak: [%3d/%3d] Loss: %f, SSIM: %f, MSSSIM: %f, Norm of Code: %f, BPP: %.2f' % (ei, args.num_epochs + args.res_epoch, val_loss, val_ssim, val_msssim, val_peanalty, val_bpp)) # bz = call('tar -jcvf ./code/code.tar.bz ./code', shell=True) # total_code_size = os.stat('./code/code.tar.bz').st_size # total_bpp = total_code_size * 8 / 24 / 768 / 512 writer.add_scalar('test/loss', val_loss, ei) writer.add_scalar('test/ssim', val_ssim, ei) writer.add_scalar('test/msssim', val_msssim, ei) writer.add_scalar('test/psnr', val_psnr, ei) writer.add_scalar('test/norm', val_peanalty, ei) writer.add_scalar('test/bpp', val_bpp, ei) # writer.add_scalar('test/total_bpp', total_bpp, ei) model.train() scheduler.step(train_loss) # save model if ei % args.save_every == args.save_every - 1: torch.save(model.state_dict(), f"./chkpt/{args.exp_name}/model.state") torch.save(optimizer.state_dict(), f"./chkpt/{args.exp_name}/opt.state") torch.save(scheduler.state_dict(), f"./chkpt/{args.exp_name}/lr.state") torch.save(Z, f"./chkpt/{args.exp_name}/Z.state") torch.save(U, f"./chkpt/{args.exp_name}/U.state") writer.close()
multichannel=True, gaussian_weights=True, ), 'piq.ssim': piq.ssim, 'kornia.SSIM-halfloss': kornia.SSIM( window_size=11, reduction='mean', ), 'piq.SSIM-loss': piq.SSIMLoss(), 'IQA.SSIM-loss': IQA.SSIM(), 'vainf.SSIM': vainf.SSIM(data_range=1.), 'piqa.SSIM': piqa.SSIM(), }), 'MS-SSIM': (2, { 'piq.ms_ssim': piq.multi_scale_ssim, 'piq.MS_SSIM-loss': piq.MultiScaleSSIMLoss(), 'IQA.MS_SSIM-loss': IQA.MS_SSIM(), 'vainf.MS_SSIM': vainf.MS_SSIM(data_range=1.), 'piqa.MS_SSIM': piqa.MS_SSIM(), }), 'LPIPS': ( 2, { 'piq.LPIPS': piq.LPIPS(), # 'IQA.LPIPS': IQA.LPIPSvgg(),
import torch import torch.nn as nn import numpy as np from model import CAEP from utils import Kodak, GeneralDS, compute_bpp from torchvision.utils import make_grid, save_image from torch.utils.data import DataLoader import pytorch_msssim import time from torch.utils.data.sampler import SequentialSampler SSIM = pytorch_msssim.SSIM().cuda() print('Number of GPUs available: ' + str(torch.cuda.device_count())) model_before = nn.DataParallel(CAEP(15, 16).cuda()) pretrained_state_dict = torch.load( f"./chkpt/bpp_10_before_pruning/model.state") current_state_dict = model_before.state_dict() current_state_dict.update(pretrained_state_dict) model_before.load_state_dict(current_state_dict) model_before.eval() print('Done Setup Model_before_pruning.') model_after = nn.DataParallel(CAEP(15, 16).cuda()) pretrained_state_dict = torch.load(f"./chkpt/bpp_05/model.state") #print(pretrained_state_dict) current_state_dict = model_after.state_dict() current_state_dict.update(pretrained_state_dict) model_after.load_state_dict(current_state_dict) model_after.eval()
def __init__(self, **kwargs): super().__init__(**kwargs) self.ssim = ssim.SSIM(data_range=1)
def main(): cudnn.benchmark = True # Dataset train_data = DatasetFromHdf5('./Data/train_Material_.h5') print(len(train_data)) val_data = DatasetFromHdf5('./Data/valid_Material_.h5') print(len(val_data)) # Data Loader (Input Pipeline) train_data_loader = DataLoader(dataset=train_data, num_workers=1, batch_size=64, shuffle=True, pin_memory=True) val_loader = DataLoader(dataset=val_data, num_workers=1, batch_size=1, shuffle=False, pin_memory=True) # Model model = resblock(conv_bn_relu_res_block, 10, 25, 25) if torch.cuda.device_count() > 1: model = nn.DataParallel(model) if torch.cuda.is_available(): model.cuda() # Parameters, Loss and Optimizer start_epoch = 0 end_epoch = 100 init_lr = 0.0001 iteration = 0 record_test_loss = 1000 # criterion_RRMSE = torch.nn.L1Loss() criterion_RRMSE = rrmse_loss criterion_Angle = Angle_Loss criterion_MSE = torch.nn.MSELoss() criterion_SSIM = pytorch_msssim.SSIM() # criterion_Div = Divergence_Loss criterion_Div = torch.nn.KLDivLoss() optimizer = torch.optim.Adam(model.parameters(), lr=init_lr, betas=(0.9, 0.999), eps=1e-08, weight_decay=0.01) model_path = './models/' if not os.path.exists(model_path): os.makedirs(model_path) loss_csv = open(os.path.join(model_path, 'loss_material.csv'), 'w+') log_dir = os.path.join(model_path, 'train_material.log') logger = initialize_logger(log_dir) # Resume resume_file = '' if resume_file: if os.path.isfile(resume_file): print("=> loading checkpoint '{}'".format(resume_file)) checkpoint = torch.load(resume_file) start_epoch = checkpoint['epoch'] iteration = checkpoint['iter'] model.load_state_dict(checkpoint['state_dict']) optimizer.load_state_dict(checkpoint['optimizer']) for epoch in range(start_epoch + 1, end_epoch): start_time = time.time() train_loss, iteration, lr = train(train_data_loader, model, criterion_MSE, criterion_RRMSE, criterion_Angle, criterion_SSIM, criterion_Div, optimizer, iteration, init_lr, end_epoch, epoch) test_loss, loss_angle, loss_reconstruct, loss_SSIM, loss_Div = validate( val_loader, model, criterion_MSE, criterion_RRMSE, criterion_Angle, criterion_SSIM, criterion_Div) # xxx_loss = validate_save(val_loader, model, criterion_MSE, criterion_RRMSE, epoch) save_checkpoint_material(model_path, epoch, iteration, model, optimizer) # print loss end_time = time.time() epoch_time = end_time - start_time print( "Epoch [%d], Iter[%d], Time:%.9f, learning rate : %.9f, Train Loss: %.9f Test Loss: %.9f , Angle Loss: %.9f, Recon Loss: %.9f, SSIM Loss: %.9f , Div Loss: %.9f" % (epoch, iteration, epoch_time, lr, train_loss, test_loss, loss_angle, loss_reconstruct, loss_SSIM, loss_Div)) # save loss record_loss(loss_csv, epoch, iteration, epoch_time, lr, train_loss, test_loss) logger.info( "Epoch [%d], Iter[%d], Time:%.9f, learning rate : %.9f, Train Loss: %.9f Test Loss: %.9f, Angle Loss: %.9f, Recon Loss: %.9f, SSIM Loss: %.9f, Div Loss: %.9f " % (epoch, iteration, epoch_time, lr, train_loss, test_loss, loss_angle, loss_reconstruct, loss_SSIM, loss_Div))
################################################################################ ##################################### Model #################################### ################################################################################ # Create Enhancer model = cnn_enhancer.ImageEnhancerCNN(input_channels, num_filters, num_layers, use_4by4_conv) model.to(device) print(model) optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) ssim_module = None if use_ssim_loss: ssim_module = pytorch_msssim.SSIM(data_range=1.0, win_size=11, win_sigma=1.5, K=(0.01, 0.03)) ################################################################################ ################################### Training ################################### ################################################################################ # Train all_train_loss = [] all_val_loss = [] for epoch in range(epochs): train_loss = 0 val_loss = 0 # Training Loop model.train()
npImg1 = cv2.imread("einstein.png") img1 = torch.from_numpy(np.rollaxis(npImg1, 2)).float().unsqueeze(0) / 255.0 img2 = torch.rand(img1.size()) if torch.cuda.is_available(): img1 = img1.cuda() img2 = img2.cuda() img1 = Variable(img1, requires_grad=False) img2 = Variable(img2, requires_grad=True) # Functional: pytorch_msssim.msssim(img1, img2, window_size = 11, size_average = True) msssim_value = pytorch_msssim.msssim(img1, img2) print("Initial msssim:", msssim_value) # Module: pytorch_msssim.SSIM(window_size = 11, size_average = True) msssim_loss = pytorch_msssim.SSIM() optimizer = optim.Adam([img2], lr=0.01) msssim_out = msssim_loss(img1, img2) ''' while msssim_value < 0.95: optimizer.zero_grad() msssim_out = -msssim_loss(img1, img2) msssim_value = - msssim_out.data[0] print(msssim_value) msssim_out.backward() optimizer.step() '''
def forward(self, flabel, image, ds_image, infer=False): # Encode Inputs input_flabel, real_image, ds_image = self.encode_input(flabel, image, ds_image) # input to G: downsampled | compact image if self.opt.comp_type=='compG': # use compG if not self.opt.no_seg: # with seg compG_input = torch.cat((input_flabel, real_image), dim=1) else: # no seg compG_input = real_image; comp_image = self.compG.forward(compG_input) ### tensor-level bilinear upsample = torch.nn.Upsample(scale_factor=self.opt.alpha, mode='bilinear') up_image = upsample(comp_image) else: # use bicubic downsampling (ds) up_image = ds_image if not self.opt.no_seg: # seg if self.opt.comp_type!='none': # seg, ds | comp_image input_fconcat = torch.cat((input_flabel, up_image), dim=1) else: # no ds, but seg input_fconcat = input_flabel else: # no seg input_flabel = None if self.opt.comp_type != 'none': # ds (ds | comp_image) input_fconcat = up_image # add compact image, so that G tries to find the best residual res = self.netG.forward(input_fconcat) fake_image_f = res + up_image # Fake Detection and Loss pred_fake_pool_f = self.discriminate(input_flabel, fake_image_f, use_pool=True) loss_D_fake = self.criterionGAN(pred_fake_pool_f, False) # Real Detection and Loss pred_real_f = self.discriminate(input_flabel, real_image) loss_D_real = self.criterionGAN(pred_real_f, True) # GAN loss (Fake Passability Loss) if input_flabel is not None: inputD_concat = torch.cat((input_flabel, fake_image_f), dim=1) else: inputD_concat = fake_image_f pred_fake_f = self.netD.forward(inputD_concat) loss_G_GAN = self.criterionGAN(pred_fake_f, True) # GAN feature matching loss loss_G_GAN_Feat = 0 if not self.opt.no_ganFeat_loss: feat_weights = 4.0 / (self.opt.n_layers_D + 1) D_weights = 1.0 / self.opt.num_D for i in range(self.opt.num_D): for j in range(len(pred_fake_f[i])-1): loss_G_GAN_Feat += D_weights * feat_weights * \ self.criterionFeat(pred_fake_f[i][j], pred_real_f[i][j].detach()) * self.opt.lambda_feat # VGG feature matching loss loss_G_VGG = 0 if not self.opt.no_vgg_loss: loss_G_VGG = (self.criterionVGG(fake_image_f, real_image)) * self.opt.lambda_feat # l1 loss between x and x' loss_G_DIS = 0 criterionDIS = torch.nn.L1Loss() loss_G_DIS = criterionDIS(fake_image_f, real_image) * self.opt.lambda_feat * 2 # SSIM Loss loss_G_SSIM=0 ssim_loss = pytorch_msssim.SSIM() loss_G_SSIM = -ssim_loss(real_image, fake_image_f) # Only return the fake_B image if necessary to save BW return [ [ loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_G_DIS, loss_G_SSIM, loss_D_real, loss_D_fake ], None if not infer else real_image, None if not infer else input_flabel, None if not infer else fake_image_f, None if not infer else res, None if not infer else comp_image, None if not infer else up_image ]
import pandas as pd import numpy as np import torch from GAN_hessian_compute import hessian_compute # from hessian_analysis_tools import scan_hess_npz, plot_spectra, average_H, compute_hess_corr, plot_consistency_example # from hessian_axis_visualize import vis_eigen_explore, vis_eigen_action, vis_eigen_action_row, vis_eigen_explore_row from GAN_utils import loadStyleGAN2, StyleGAN2_wrapper, loadBigGAN, BigGAN_wrapper, loadPGGAN, PGGAN_wrapper import matplotlib.pylab as plt import matplotlib #%% device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') import lpips ImDist = lpips.LPIPS(net="squeeze").cuda() # SSIM import pytorch_msssim D = pytorch_msssim.SSIM( ) # note SSIM, higher the score the more similar they are. So to confirm to the distance convention, we use 1 - SSIM as a proxy to distance. # L2 / MSE def MSE(im1, im2): return (im1 - im2).pow(2).mean(dim=[1, 2, 3]) # L1 / MAE # def L1(im1, im2): # return (im1 - im2).abs().mean(dim=[1,2,3]) # Note L1 is less proper for this task, as it's farther away from a distance square function. #%% Utility functions to quantify relationship between 2 eigvals or 2 Hessians. def spectra_cmp(eigvals1, eigvals2, show=True): cc = np.corrcoef((eigvals1), (eigvals2))[0, 1] logcc = np.corrcoef(np.log10(np.abs(eigvals1) + 1E-8),
def run(config, dataset_path, IPF_path, image_list, idx, fileText, args): # Create a directory for saving adversarial images adv_path_no = 'Results/{}/{}_{}_{}/'.format(args.method, str(args.strength), args.dataset, args.adv_model) if not os.path.isdir(adv_path_no): os.makedirs(adv_path_no) # Structure loss function l2_loss = nn.MSELoss() ssim_loss = pytorch_msssim.SSIM() # Using GPU if config.GPU >= 0: with torch.cuda.device(config.GPU): config.model.cuda() l2_loss.cuda() ssim_loss.cuda() # Setup optimizer optimizer = optim.Adam(config.model.parameters(), lr=config.LR) # Load the classifier for attacking if args.adv_model == 'resnet18': classifier = models.resnet18(pretrained=True) elif args.adv_model == 'resnet50': classifier = models.resnet50(pretrained=True) elif args.adv_model == 'alexnet': classifier = models.alexnet(pretrained=True) classifier.cuda() classifier.eval() # Freeze the parameters of the classifier under attack to not be updated for param in classifier.parameters(): param.requires_grad = False # The name of the chosen image img_name = image_list[idx].split('/')[-1] # Load and Pre-processing the clean and filtered image x = processImage(dataset_path, img_name) gt_enh = processImage(IPF_path, img_name) # Compute the residual perturbation gt_noise = gt_enh - x # Perform inference on the clean image class_x, prob_class_x, prob_x, logit_x, semantic_vec, super_x = PreidictLabel( x, classifier) maxIters = 3000 for it in tqdm(range(maxIters)): with autograd.detect_anomaly(): noise = config.forward(x, gt_noise, config) # Enhance adversarial image enh = (x + noise).clamp(min=0, max=1) # Perform inference on the generated adversarial image class_enh, prob_class_enh, prob_enh, logit_enh, _, super_adv = PreidictLabel( enh, classifier) # Computing structure and semantic adversarial losses loss0 = l2_loss(noise, gt_noise) loss1 = 1 - ssim_loss((noise + 1) / 2., (gt_noise + 1) / 2.) loss2 = AdvLoss(logit_enh, class_x, semantic_vec, is_targeted=False, num_classes=1000) # Normalized MSE loss3 = loss0.cpu().data.numpy().item(0) / l2_loss( gt_noise, torch.zeros(1, 3, 224, 224).cuda()) loss = loss0 + 0.01 * loss1 + loss2 # backward optimizer.zero_grad() loss.backward() if config.clip is not None: torch.nn.utils.clip_grad_norm(config.model.parameters(), config.clip) optimizer.step() # Save the generated adversarial image cv2.imwrite('{}{}'.format(adv_path_no, img_name), recreate_image(enh)) adv_img = processImage(adv_path_no, img_name) class_adv, _, _, _, _, super_adv = PreidictLabel( adv_img, classifier) if args.method == 'Nonlinear_Detail': if (super_x != super_adv and class_x != class_adv and loss3 < 0.04 and it > 2500): break elif args.method == 'Log': if (super_x != super_adv and class_x != class_adv and loss3 < 0.003 and it > 2500): break elif args.method == 'Linear_Detail': if (super_x != super_adv and class_x != class_adv and loss3 < 0.04 and it > 2500): break elif args.method == 'Gamma': if (super_x != super_adv and class_x != class_adv and loss3 < 0.0005 and it > 2500): break #print(img_name, it+1, super_x, super_adv, class_x.cpu().data.numpy().item(0), class_enh.cpu().data.numpy().item(0), class_adv.cpu().data.numpy().item(0), loss0.cpu().data.numpy().item(0), loss3.cpu().data.numpy().item(0), loss1.cpu().data.numpy().item(0), loss2.cpu().data.numpy().item(0)) text = '{}\tItrs:{}\tSemantic labels, Clean:{}\t Adversarial:{}\t Categorical labels, Clean:{}\t Adversarial:{}\t L_2 loss:{:.5f}\t SSIM loss{:.5f}\t Adv loss:{:.5f}\n'.format( img_name, it + 1, super_x, super_adv, class_x, class_adv, loss0.cpu().data.numpy().item(0), loss1.cpu().data.numpy().item(0), loss2.cpu().data.numpy().item(0)) fileText.write(text) return adv_img