Organize the code and add comments

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/'

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..')

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

deepsecure/cat_adv_image.py

@@ -1,126 +1,17 @@
-# 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
 adv_img_path = cfg.adv_img_path
 
 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,19 +100,16 @@ 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
         """
         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")

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.

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):
 

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())

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

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