diff --git a/config.py b/config.py index 95dbb0f..8ef7aa7 100644 --- a/config.py +++ b/config.py @@ -1,24 +1,32 @@ -use_cuda = True -image_nc = 3 -epochs = 80 -batch_size = 64 -BOX_MIN = 0 +# if use GPU +use_cuda = True +image_nc = 3 +# training epoch +epochs = 80 +# batch size +batch_size = 64 +# minimum pixel intensity value +BOX_MIN = 0 +# maximum pixel intensity value +BOX_MAX = 1 -BOX_MAX = 1 pretrained_model_arch = 'resnet18' -num_layers_ext = 5 -ext_fixed = True +num_layers_ext = 5 +ext_fixed = True -G_tagged = True -tag_size = 6 +G_tagged = True +# tag size +tag_size = 6 -cat_G = True -noise_img = True -noise_g_path = './models/netG_epoch_80.pth' +cat_G = True +noise_img = True +noise_g_path = './models/netG_epoch_80.pth' -imagenet10_traindir = 'C:/Users\Rui\Pictures/transfer_imgnet_10/train' -imagenet10_valdir = 'C:/Users\Rui\Pictures/transfer_imgnet_10/val' +imagenet10_traindir = 'C:/Users\Rui\Pictures/transfer_imgnet_10/train' +imagenet10_valdir = 'C:/Users\Rui\Pictures/transfer_imgnet_10/val' + + +models_path = './models/' +adv_img_path = './images/0526/adv/' -models_path = './models/' -adv_img_path = './images/0526/adv/' diff --git a/data/imagenet10_dataloader.py b/data/imagenet10_dataloader.py index cad8037..c265120 100644 --- a/data/imagenet10_dataloader.py +++ b/data/imagenet10_dataloader.py @@ -1,7 +1,7 @@ import torch import torchvision.datasets as datasets import torchvision.transforms as transforms -import config +import config def get_data_loaders(): print('==> Preparing Imagenet 10 class data..') diff --git a/deepsecure/adv_image.py b/deepsecure/adv_image.py index 42869aa..0fa6913 100644 --- a/deepsecure/adv_image.py +++ b/deepsecure/adv_image.py @@ -4,7 +4,7 @@ import torch.nn.functional as F import torchvision import os -import config as cfg +import config as cfg models_path = cfg.models_path adv_img_path = cfg.adv_img_path diff --git a/deepsecure/cat_adv_image.py b/deepsecure/cat_adv_image.py index d0f2808..dd6e2da 100644 --- a/deepsecure/cat_adv_image.py +++ b/deepsecure/cat_adv_image.py @@ -1,118 +1,10 @@ -# import torch.nn as nn -# import torch -# import torch.nn.functional as F -# import torchvision -# import os -# import config as cfg - -# models_path = cfg.models_path -# adv_img_path = cfg.adv_img_path - -# # custom weights initialization called on netG and netD -# def weights_init(m): -# classname = m.__class__.__name__ -# if classname.find('Conv') != -1: -# nn.init.normal_(m.weight.data, 0.0, 0.02) -# elif classname.find('BatchNorm') != -1: -# nn.init.normal_(m.weight.data, 1.0, 0.02) -# nn.init.constant_(m.bias.data, 0) - - -# class Cat_Adv_Gen: -# def __init__(self, -# device, -# model_extractor, -# generator, -# reg_g): - -# self.device = device -# self.model_extractor = model_extractor -# self.generator = generator -# self.box_min = cfg.BOX_MIN -# self.box_max = cfg.BOX_MAX -# self.ite = 0 -# #self.CELoss = nn.CrossEntropyLoss() - -# self.model_extractor.to(device) -# #self.model_extractor.eval() - -# self.generator.to(device) - -# self.noise_generator = reg_g -# if self.noise_generator != False: -# self.noise_generator.to(device) - -# # initialize optimizers -# self.optimizer_G = torch.optim.Adam(self.generator.parameters(), -# lr=0.001) - -# if not os.path.exists(models_path): -# os.makedirs(models_path) -# if not os.path.exists(adv_img_path): -# os.makedirs(adv_img_path) - -# def train_batch(self, x): -# self.optimizer_G.zero_grad() -# idx = torch.randint(x.size(0),(x.size(0),)) -# if self.noise_generator != False: -# with torch.no_grad(): -# noise_imgs, temp_feature = self.noise_generator(x) -# adv_imgs, hide_feature = self.generator(x, noise_imgs[idx]) -# else: -# adv_imgs, hide_feature = self.generator(x, x[idx]) -# adv_img_feature = self.model_extractor(adv_imgs) - -# loss_img = F.l1_loss(adv_imgs, x+5*noise_imgs) -# loss_adv = F.l1_loss(hide_feature, adv_img_feature) -# loss_G = loss_img + 1*loss_adv -# loss_G.backward(retain_graph=True) - -# self.optimizer_G.step() - -# return loss_adv.item(), adv_imgs, idx, loss_img.item() - -# def train(self, train_dataloader, epochs): -# for epoch in range(1, epochs+1): - -# if epoch == 200: -# self.optimizer_G = torch.optim.Adam(self.netG.parameters(), -# lr=0.0001) -# if epoch == 400: -# self.optimizer_G = torch.optim.Adam(self.netG.parameters(), -# lr=0.00001) -# loss_adv_sum = 0 -# loss_img_sum = 0 -# self.ite = epoch -# for i, data in enumerate(train_dataloader, start=0): -# images, labels = data -# images, labels = images.to(self.device), labels.to(self.device) - -# loss_adv_batch, adv_img, idx, loss_img_batch = self.train_batch(images) -# loss_adv_sum += loss_adv_batch -# loss_img_sum += loss_img_batch - - -# # print statistics - -# torchvision.utils.save_image(torch.cat((adv_img[:7], images[:7], (images[idx])[:7])), -# adv_img_path + str(epoch) + ".png", -# normalize=True, scale_each=True, nrow=7) -# num_batch = len(train_dataloader) -# print("epoch %d:\n loss_adv: %.3f, loss_img: %.3f \n" % -# (epoch, loss_adv_sum/num_batch, loss_img_sum/num_batch)) -# # save generator -# if epoch%20==0: -# netG_file_name = models_path + 'catted_netG_epoch_' + str(epoch) + '.pth' -# torch.save(self.generator.state_dict(), netG_file_name) - -# print("check") -#import +import os import torch.nn as nn import torch -import torch.nn.functional as F import torchvision -import os import config as cfg +import torch.nn as nn +import torch.nn.functional as F # Get paths from config for saving models and adversarial images models_path = cfg.models_path @@ -120,7 +12,6 @@ def weights_init(m): """Custom weight initialization for the generator and discriminator networks. - Args: m: Module whose weights need to be initialized """ @@ -135,7 +26,6 @@ def weights_init(m): class Cat_Adv_Gen: """Concatenated Adversarial Generator class that combines regular and noise generators.""" - def __init__(self, device, model_extractor, generator, reg_g): """Initialize the Cat_Adv_Gen with necessary components. @@ -162,8 +52,7 @@ def __init__(self, device, model_extractor, generator, reg_g): self.noise_generator.to(device) # Initialize Adam optimizer for generator - self.optimizer_G = torch.optim.Adam(self.generator.parameters(), - lr=0.001) + self.optimizer_G = torch.optim.Adam(self.generator.parameters(), lr=0.001) # Create directories for saving models and images if they don't exist if not os.path.exists(models_path): @@ -173,7 +62,6 @@ def __init__(self, device, model_extractor, generator, reg_g): def train_batch(self, x): """Train generator on a single batch of images. - Args: x: Batch of input images @@ -212,7 +100,6 @@ def train_batch(self, x): def train(self, train_dataloader, epochs): """Train the generator for specified number of epochs. - Args: train_dataloader: DataLoader containing training images epochs: Number of epochs to train for @@ -220,11 +107,9 @@ def train(self, train_dataloader, epochs): for epoch in range(1, epochs+1): # Reduce learning rate at specific epochs if epoch == 200: - self.optimizer_G = torch.optim.Adam(self.netG.parameters(), - lr=0.0001) + self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=0.0001) if epoch == 400: - self.optimizer_G = torch.optim.Adam(self.netG.parameters(), - lr=0.00001) + self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=0.00001) # Initialize loss sums for epoch loss_adv_sum = 0 @@ -255,5 +140,4 @@ def train(self, train_dataloader, epochs): if epoch % 20 == 0: netG_file_name = models_path + 'catted_netG_epoch_' + str(epoch) + '.pth' torch.save(self.generator.state_dict(), netG_file_name) - #print print("check") diff --git a/deepsecure/catted_generator.py b/deepsecure/catted_generator.py index 7c7d02e..50cc1c3 100644 --- a/deepsecure/catted_generator.py +++ b/deepsecure/catted_generator.py @@ -71,8 +71,7 @@ import torch.nn as nn import torch from module.resnet_block import ResnetBlock -from module.pre_model_extractor import model_extractor -import config as cfg +from module.pre_model_extractor import model_extractor class catted_generator(nn.Module): """Concatenated Generator that processes and combines features from two inputs. diff --git a/deepsecure/regular_generator.py b/deepsecure/regular_generator.py index ae09a22..eb91c9c 100644 --- a/deepsecure/regular_generator.py +++ b/deepsecure/regular_generator.py @@ -1,9 +1,8 @@ -#import import torch.nn as nn import torchvision.models as pre_models from module.resnet_block import ResnetBlock from module.pre_model_extractor import model_extractor -import config as cfg +import config as cfg class regular_generator(nn.Module): diff --git a/main.py b/main.py index 03fd56d..205d229 100644 --- a/main.py +++ b/main.py @@ -6,8 +6,8 @@ from regular_generator import regular_generator from catted_generator import catted_generator -from pre_model_extractor import model_extractor -#------------------------------- +from pre_model_extractor import model_extractor + if __name__ == '__main__': # Check if CUDA is available print("CUDA Available: ", torch.cuda.is_available()) diff --git a/module/pre_model_extractor.py b/module/pre_model_extractor.py index 9a38007..22ec690 100644 --- a/module/pre_model_extractor.py +++ b/module/pre_model_extractor.py @@ -1,5 +1,5 @@ import torch.nn as nn -import torchvision.models as pre_models +import torchvision.models as pre_models # Return first n layers of a pretrained model diff --git a/module/resnet_block.py b/module/resnet_block.py index b8e51e3..eb4a07a 100644 --- a/module/resnet_block.py +++ b/module/resnet_block.py @@ -1,47 +1,5 @@ -# import torch.nn as nn - -# # Define a resnet block -# # modified from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py -# class ResnetBlock(nn.Module): -# def __init__(self, dim, padding_type='reflect', norm_layer=nn.BatchNorm2d, use_dropout=False, use_bias=False): -# super(ResnetBlock, self).__init__() -# self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias) - -# def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias): -# conv_block = [] -# p = 0 -# if padding_type == 'reflect': -# conv_block += [nn.ReflectionPad2d(1)] -# elif padding_type == 'replicate': -# conv_block += [nn.ReplicationPad2d(1)] -# elif padding_type == 'zero': -# p = 1 -# else: -# raise NotImplementedError('padding [%s] is not implemented' % padding_type) - -# conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), -# norm_layer(dim), -# nn.ReLU(True)] -# if use_dropout: -# conv_block += [nn.Dropout(0.5)] - -# p = 0 -# if padding_type == 'reflect': -# conv_block += [nn.ReflectionPad2d(1)] -# elif padding_type == 'replicate': -# conv_block += [nn.ReplicationPad2d(1)] -# elif padding_type == 'zero': -# p = 1 -# else: -# raise NotImplementedError('padding [%s] is not implemented' % padding_type) - -# conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), -# norm_layer(dim)] - -# return nn.Sequential(*conv_block) import torch.nn as nn - # Define a resnet block # modified from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py