Saved autoencoder gets trained repeatedly on runs with same parameters

[mnvt]

Hello, When running DeepDIG with the same parameters as a previous run (i.e. same dataset, DNN type and classes), I noticed that every subsequent run loads the autoencoder trained in the previous run, but then trains it a second time for, again, 5000 iterations. I would have expected that loading an AE from a previous run with identical parameters would cause the AE training portion to be skipped and the existing AE to simply be re-used.

Here is the first part (only showing class 1-->2) of the output of my initial run:

$ python3 -m DeepDIGCode.main --dataset MNIST --pre-trained-model CNN --classes "1;2"
Running DeepDIG for 1 --> 2 ...
 No trained model, creating one ...
Step 0 Loss 0.9017552137374878
Step 2000 Loss 0.0667022317647934
Step 4000 Loss 0.0745028704404831
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv/model.m loaded
6741/6742 1->2 adv samples successfully generated 
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv/model.m loaded
1130/1135 1->2 adv samples successfully generated 
 No trained model, creating one ...
Step 0 Loss 0.7947918176651001
Step 2000 Loss 0.06338567286729813
Step 4000 Loss 0.0638570562005043
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv_of_adv/model.m loaded
6741/6741 1->2->1 adv_of_adv samples successfully generated
6741it [01:42, 65.85it/s]
/home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/borderline_deepdig_0.0001.pkl saved
Class 1->2 
    success rate 1.0
    avg attempts success 14.971665924936953
    avg attempts fail -1
    avg attempts total 14.971665924936953
****************************************************************************************************
Finished DeepDIG for 1 --> 2
[...]

As expected, no existing AE was found, so a new one was created. Yet, when running it a second time, I get this:

$ python3 -m DeepDIGCode.main --dataset MNIST --pre-trained-model CNN --classes "1;2"
Running DeepDIG for 1 --> 2 ...
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv/model.m loaded
Step 0 Loss 0.06560038030147552
Step 2000 Loss 0.06348234415054321
Step 4000 Loss 0.059069789946079254
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv/model.m loaded
6739/6742 1->2 adv samples successfully generated 
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv/model.m loaded
1130/1135 1->2 adv samples successfully generated 
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv_of_adv/model.m loaded
Step 0 Loss 0.06379116326570511
Step 2000 Loss 0.06314456462860107
Step 4000 Loss 0.06670735031366348
Saved model /home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/adv_of_adv/model.m loaded
6739/6739 1->2->1 adv_of_adv samples successfully generated
6739it [01:47, 62.67it/s]
/home/user/Downloads/DeepDIGCode/PreTrainedModels/MNIST/CNN/(1,2)/1_2/borderline_deepdig_0.0001.pkl saved
Class 1->2 
    success rate 0.998664490280457
    avg attempts success 15.254531946508173
    avg attempts fail 2.4444444444444446
    avg attempts total 15.23742395014097
****************************************************************************************************
Finished DeepDIG for 1 --> 2
[...]

The train() functions in train_adv.py and train_adv_of_adv.py load the saved AE from the previous run, but then train it for another 5000 steps (where the previous run left off) instead of skipping that part and leaving the pre-trained AE as is because the training code is not indented to be in the else-block (seen below). I wanted to ask if this is intentional? And if it is, why is the existing AE trained for an additional args.steps?

def train():
    model = utils.ae_module.AE(dropout=args.dropout).to(utils.device)
    if os.path.exists(save_dir+'model.m'):
        model = torch.load(save_dir+'model.m')
        print("Saved model {} loaded".format(save_dir+'model.m'))
    else:
        print(" No trained model, creating one ...")
    optimizer = optim.Adam(model.parameters(), lr=args.lr)

    scheduler = StepLR(optimizer, step_size=args.step_size_scheduler, gamma=args.gamma_scheduler)

    for step in range(args.steps):
        perm = torch.randperm(len(train_data))
        batch = train_data[perm[0:args.batch_size]].float().to(utils.device)
        Z,_ = pre_trained_model(batch.view(utils.get_pre_trained_model_input_shape()))
        pre_trained_model_data_probs = F.softmax(Z, dim=1)

        _,recon_batch = model(batch)

        Z,_= pre_trained_model(recon_batch.view(utils.get_pre_trained_model_input_shape()).to(utils.device))
        pre_trained_model_reconst_probs = F.softmax(Z, dim=1)
        loss = utils.loss_function_ae(recon_batch, batch, pre_trained_model_reconst_probs,
                                          pre_trained_model_data_probs, args.batch_size,
                                          utils.classes['t'], device=utils.device)
        optimizer.zero_grad()
        loss.backward()

        optimizer.step()

        scheduler.step()

        if step % 2000==0:
            print("Step {} Loss {}".format(step,loss.cpu().data))
    torch.save(model,save_dir+'model.m')

[hamidkarimi]

@mnvt A good question. Yes, you can modify the code and add a simple return after a saved model is found:

if os.path.exists(save_dir+'model.m'):
        model = torch.load(save_dir+'model.m')
        print("Saved model {} loaded".format(save_dir+'model.m'))
        return

At the time, my intention was by calling train(), the user wants to train the AE model anyway. If such a model doesn't exist, the new one will be trained. Otherwise, the pre-trained one is loaded and re-trained again. For instance, for a certain dataset, a user just trains AE model for say 1000 steps and then later decides to train it more.

[mnvt]

Thanks for the answer.

Additionally, I noticed that you used the binary cross entropy loss for both the left and the right part of the autoencoder loss function, when the paper suggests a squared L2 distance for the reconstruction error. So I again wanted to ask the motivation behind this? I suppose a better performance (higher adversarial example generation success rate/ adversarial example closer to original sample) when all other simulation parameters stay the same?

reconst_loss_function = torch.nn.BCELoss()
BCELoss = torch.nn.BCELoss()

def loss_function_ae(reconst_ae, data, pre_trained_model_reconst_probs,
                     pre_trained_model_data_probs, batch_size, target, device):
    prediction_pre_trained_model = pre_trained_model_data_probs.max(1, keepdim=True)[1].cpu().numpy().reshape(batch_size)
    reconstrtion_loss = reconst_loss_function(reconst_ae, data)

    attacked_targets = torch.zeros(pre_trained_model_data_probs.shape)
    for i, j in zip(range(attacked_targets.shape[0]), prediction_pre_trained_model):
        attacked_targets[i][target] = 1

    attacked_targets = attacked_targets.to(device)

    targte_loss = BCELoss(pre_trained_model_reconst_probs, attacked_targets)
    return reconstrtion_loss  + args.alpha * targte_loss