Fix training and experiments

Author: FrederickGeek8

experiments/test.py

@@ -33,7 +33,8 @@ def parse_args():
 
 
 if __name__ == "__main__":    
-    models = ['vanilla', 'classification', 'proxi_dist', 'combined']
+    # models = ['classification', 'proxi_dist', 'vanilla', 'combined']
+    models = ['classification', 'vanilla']
     for i in range(len(models)):
         args = parse_args()
         model = MADVAE(args)

experiments/test_black.py

@@ -5,10 +5,9 @@
 from torch.utils import data
 from torchvision import datasets, transforms
 import matplotlib.pyplot as plt
+from test.test_models import *
 sys.path.insert(0, os.path.abspath('..'))
 from MAD_VAE import *
-from test.test_models import *
-from test.plotting import *
 from utils.dataset import *
 from utils.adversarial import *
 from utils.classifier import *
@@ -35,7 +34,7 @@ def parse_args():
 
 if __name__ == "__main__":
     models = ['vanilla', 'classification', 'proxi_dist', 'combined']
-    for i in range(4):
+    for i in range(len(models)):
 
         args = parse_args()
         model = MADVAE(args)
@@ -84,23 +83,23 @@ def parse_args():
                 image = image.cuda()
                 label = label.cuda()
 
-                output, adv_out = add_adv(classifier, image, label, 'fgsm')
+                output, adv_out = add_adv(classifier, image, label, 'fgsm', default=True)
                 output_class = classifier(output)
                 adv_output_class = classifier(adv_out)
                 def_out, _, _, _ = model(adv_out)
-                adv_out_class = classifier(def_out)
+                cleaned_class = classifier(def_out)
 
                 true_class = torch.argmax(output_class, 1)
-                output_class = torch.argmax(adv_output_class, 1)
-                adversarial_class = torch.argmax(adv_out_class, 1)
+                adv_class = torch.argmax(adv_output_class, 1)
+                adv_clean_class = torch.argmax(cleaned_class, 1)
 
                 print(f'attack method fgsm')
                 print(f'actual class {true_class}')
-                print(f'actual advclass {output_class}')
-                print(f'adversarial class {adversarial_class}')
+                print(f'actual advclass {adv_class}')
+                print(f'adversarial class {adv_clean_class}')
 
-                true += torch.sum(torch.eq(true_class, adversarial_class))
-                true_adv += torch.sum(torch.eq(true_class, output_class))
+                true += torch.sum(torch.eq(true_class, adv_clean_class))
+                true_adv += torch.sum(torch.eq(true_class, adv_class))
 
                 print(int(true) / total)
                 print(int(true_adv) / total)

pretrained_model/classification/params.pt

@@ -111,7 +111,7 @@ def train(args, dataloader, model, optimizer, step, epoch):
         distribution = Normal(dsm, dss)
 
         # calculate losses
-        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, epoch/100)
+        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, 0.1)
         loss = r_loss
         loss.backward()
 

pretrained_model/combined/params.pt

@@ -116,7 +116,7 @@ def train(args, dataloader, model, classifier, optimizer, step, epoch):
         distribution = Normal(dsm, dss)
 
         # calculate losses
-        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, epoch/100)
+        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, 0.1)
         c_loss = classification_loss(output, label, classifier)
         loss = r_loss + args.closs_weight * c_loss
         loss.backward()

pretrained_model/proxi_dist/params.pt

@@ -18,17 +18,16 @@ def parse_args():
     desc = "MAD-VAE for adversarial defense"
     parser = argparse.ArgumentParser(description=desc)
     parser.add_argument('--batch_size', type=int, default=512, help='Training batch size')
-    parser.add_argument('--epochs', type=int, default=10, help='Training epoch numbers')
+    parser.add_argument('--epochs', type=int, default=5, help='Training epoch numbers')
     parser.add_argument('--h_dim', type=int, default=4096, help='Hidden dimensions')
     parser.add_argument('--z_dim', type=int, default=128, help='Latent dimensions for images')
     parser.add_argument('--image_channels', type=int, default=1, help='Image channels')
     parser.add_argument('--image_size', type=int, default=28, help='Image size (default to be squared images)')
     parser.add_argument('--num_classes', type=int, default=10, help='Number of image classes')
     parser.add_argument('--log_dir', type=str, default='pd_logs', help='Logs directory')
     parser.add_argument('--lr', type=float, default=0.001, help='Learning rate for the Adam optimizer')
-    parser.add_argument('--closs_weight', type=float, default=0.1, help='Weight for classification loss functions')
     parser.add_argument('--ploss_weight', type=float, default=0.01, help='Weight for proximity loss functions')
-    parser.add_argument('--dloss_weight', type=float, default=0.00001, help='Weight for distance loss functions')
+    parser.add_argument('--dloss_weight', type=float, default=0.0001, help='Weight for distance loss functions')
     parser.add_argument('--data_root', type=str, default='data', help='Data directory')
     parser.add_argument('--model_dir', type=str, default='pretrained_model', help='Pretrained model directory')
     parser.add_argument('--use_gpu', type=bool, default=True, help='If use GPU for training')
@@ -78,18 +77,18 @@ def main():
             writer1.add_image("reconstruct data", outputs[i][0], step)
 
         # print out loss
-        print("batch {}'s img_recon loss: {:.5f}, recon loss: {:.5f}, kl loss: {:.5f}"\
+        print("batch {}'s img_recon loss: {:.5f}, recon loss: {:.5f}, kl loss: {:.5f}, pd_loss: {:.5f}"\
             .format(step, np.sum(img_losses)/len(img_losses), np.sum(recon_losses)/len(recon_losses),\
-                    np.sum(kl_losses)/len(kl_losses)))
+                    np.sum(kl_losses)/len(kl_losses), np.sum(pd_losses)/len(pd_losses)))
 
         # step scheduler
         scheduler.step()
         scheduler1.step()
         scheduler2.step()
 
         # save model parameters
-        if epoch % 5 == 0:
-            torch.save(model.state_dict(), '{}/proxi_dist/params_{}.pt'.format(args.model_dir, epoch))
+        # if epoch % 5 == 0:
+            # torch.save(model.state_dict(), '{}/proxi_dist/params_{}.pt'.format(args.model_dir, epoch))
 
     torch.save(model.state_dict(), '{}/proxi_dist/params.pt'.format(args.model_dir))
 
@@ -123,10 +122,12 @@ def train(args, dataloader, model, classifier, proximity, distance, optimizer, o
         distribution = Normal(dsm, dss)
 
         # calculate losses
-        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, epoch/100)
+        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, 0.1)
         p_loss = proximity(z, label)
         d_loss = distance(z, label) 
-        loss = r_loss + + args.ploss_weight * p_loss - args.dloss_weight * d_loss
+
+        pd_loss = args.ploss_weight * p_loss - args.dloss_weight * d_loss
+        loss = r_loss + pd_loss
         loss.backward()
 
         # clip for gradient
@@ -136,18 +137,14 @@ def train(args, dataloader, model, classifier, proximity, distance, optimizer, o
 
         # step optimizer
         optimizer.step()
-        for param in proximity.parameters():
-            param.grad.data *= (1. / args.ploss_weight)
         optimizer1.step()
-        for param in distance.parameters():
-            param.grad.data *= (1. / args.dloss_weight)
         optimizer2.step()
 
         # record results
         recon_losses.append(loss.cpu().item())
         img_losses.append(img_recon.cpu().item())
         kl_losses.append(kld.cpu().item())
-        pd_losses.append(p_loss.cpu().item() - d_loss.cpu().item())
+        pd_losses.append(pd_loss)
         outputs.append(output.cpu())
         datas.append(data.cpu())
         adv_datas.append(adv_data.cpu())
@@ -191,10 +188,10 @@ def init_models(args):
     # construct optimizer
     optimizer = optim.Adam(model.parameters(), lr=args.lr)
     scheduler = MinExponentialLR(optimizer, gamma=0.998, minimum=1e-5)
-    optimizer1 = optim.Adam(proximity.parameters(), lr=args.lr*50)
-    scheduler1 = MinExponentialLR(optimizer1, gamma=0.998, minimum=1e-5)
-    optimizer2 = optim.Adam(distance.parameters(), lr=args.lr/100)
-    scheduler2 = MinExponentialLR(optimizer2, gamma=0.998, minimum=1e-5)
+    optimizer1 = optim.SGD(proximity.parameters(), lr=args.lr*500)
+    scheduler1 = MinExponentialLR(optimizer1, gamma=0.1, minimum=1e-5)
+    optimizer2 = optim.SGD(distance.parameters(), lr=args.lr/10)
+    scheduler2 = MinExponentialLR(optimizer2, gamma=0.1, minimum=1e-5)
 
     return model, proximity, distance, classifier, optimizer, scheduler,\
          optimizer1, scheduler1, optimizer2, scheduler2

pretrained_model/vanilla/params.pt

@@ -126,11 +126,12 @@ def train(args, dataloader, model, classifier, proximity, distance, optimizer, o
         distribution = Normal(dsm, dss)
 
         # calculate losses
-        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, epoch/100)
+        r_loss, img_recon, kld = recon_loss_function(output, data, distribution, step, 0.1)
         c_loss = classification_loss(output, label, classifier)
         p_loss = proximity(z, label)
         d_loss = distance(z, label)
-        loss = r_loss + args.closs_weight * c_loss + args.ploss_weight * p_loss - args.dloss_weight * d_loss
+        pd_loss = args.ploss_weight * p_loss - args.dloss_weight * d_loss
+        loss = r_loss + args.closs_weight * c_loss + pd_loss
         loss.backward()
 
         # clip for gradient
@@ -140,19 +141,15 @@ def train(args, dataloader, model, classifier, proximity, distance, optimizer, o
 
         # step optimizer
         optimizer.step()
-        for param in proximity.parameters():
-            param.grad.data *= (1. / args.ploss_weight)
         optimizer1.step()
-        for param in distance.parameters():
-            param.grad.data *= (1. / args.dloss_weight)
         optimizer2.step()
 
         # record results
         recon_losses.append(loss.cpu().item())
         img_losses.append(img_recon.cpu().item())
         kl_losses.append(kld.cpu().item())
         c_losses.append(c_loss.cpu().item())
-        pd_losses.append(p_loss.cpu().item() - d_loss.cpu().item())
+        pd_losses.append(pd_loss)
         outputs.append(output.cpu())
         datas.append(data.cpu())
         adv_datas.append(adv_data.cpu())
@@ -196,10 +193,10 @@ def init_models(args):
     # construct optimizer
     optimizer = optim.Adam(model.parameters(), lr=args.lr)
     scheduler = MinExponentialLR(optimizer, gamma=0.998, minimum=1e-5)
-    optimizer1 = optim.Adam(proximity.parameters(), lr=args.lr*50)
-    scheduler1 = MinExponentialLR(optimizer1, gamma=0.998, minimum=1e-5)
-    optimizer2 = optim.Adam(distance.parameters(), lr=args.lr/100)
-    scheduler2 = MinExponentialLR(optimizer2, gamma=0.998, minimum=1e-5)
+    optimizer1 = optim.SGD(proximity.parameters(), lr=args.lr*500)
+    scheduler1 = MinExponentialLR(optimizer1, gamma=0.1, minimum=1e-5)
+    optimizer2 = optim.SGD(distance.parameters(), lr=args.lr/10)
+    scheduler2 = MinExponentialLR(optimizer2, gamma=0.1, minimum=1e-5)
 
     return model, proximity, distance, classifier, optimizer, scheduler,\
          optimizer1, scheduler1, optimizer2, scheduler2