Auxiliary batch norm for AdvProp

ialhashim: <p>I am trying to implement the paper <a href="https://arxiv.org/pdf/1911.09665.pdf" rel="nofollow noopener">Adversarial Examples Improve Image Recognition</a>. A main component is having branching where some samples go through one branch and others go through another.</p>
<p>What I have so far is something like this, starting from an existing model I first replace the BN layers as such:</p>
<pre><code class="lang-auto">class MyBatchNorm2d(nn.Module):
    def __init__(self, init_bn):
        super(MyBatchNorm2d, self).__init__()
        self.bn = copy.deepcopy(init_bn)
        
        # Aux BN
        self.bn_aux = copy.deepcopy(init_bn)
        self.bn_aux.reset_parameters()
        
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        
    def forward(self, input):            
        if self.training:
            self.num_batches_tracked = self.num_batches_tracked + 1            
        
        # Odd runs are AUX input
        if self.training and self.num_batches_tracked % 2 != 0:
            return self.bn_aux(input)
      
        return self.bn(input)

def convert_bn_to_bnaux(model):
    global num_batches_tracked
    for child_name, child in model.named_children():
        if isinstance(child, nn.BatchNorm2d):
            setattr(model, child_name, MyBatchNorm2d(copy.deepcopy(child)))
            del child
        else:
            convert_bn_to_bnaux(child)
</code></pre>
<p>Then while training I do the following:</p>
<pre><code class="lang-auto">    for i, sample_batched in enumerate(train_loader):
        # Switch to train mode
        model.train()

        optimizer.zero_grad()

        # Predict
        output = model(image)
        output_aux = model(adv_image) # some adversarial image

        # Compute the loss
        loss = l1_criterion(output, labels)

        # AUX
        loss_aux = l1_criterion(output_aux, labels)

        loss_both = loss + loss_aux

        # Update step
        loss_both.backward()
        optimizer.step()
</code></pre>
<p>However, I am not sure if it the back propagation is correct in this case. Do I need to branch on the backward pass?</p>
<p>I tried to diagnose it by flipping the two training steps to try and mess up the results but I still get the same test error.</p>

ptrblck: <p>The backward pass will be executed according to what Autograd collected during the forward pass, so the condition should be respected.</p>
<p>I’m not sure to understand your code completely, but it seems you are initializing <code>self.bn</code> and <code>self.bn_aux</code> with the same values? If so, switching the order wouldn’t make a difference?</p>

ialhashim: <p>Thanks for the feedback <a class="mention" href="/u/ptrblck">@ptrblck</a>!</p>
<p>For the initial values, I call this function on one of them. Does it not zero out the BN parameters.</p>
<pre><code class="lang-auto">       self.bn_aux.reset_parameters()
</code></pre>

ptrblck: <p>Yes, you are right and I have assumed you are passing an already reset bn layer into your custom implementation.</p>
<p>Your <code>MyBatchNorm2d</code> implementation seems to work correctly:</p>
<pre><code class="lang-python">class MyBatchNorm2d(nn.Module):
    def __init__(self, init_bn):
        super(MyBatchNorm2d, self).__init__()
        self.bn = copy.deepcopy(init_bn)
        
        # Aux BN
        self.bn_aux = copy.deepcopy(init_bn)
        self.bn_aux.reset_parameters()
        
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        
    def forward(self, input):            
        if self.training:
            self.num_batches_tracked = self.num_batches_tracked + 1            
        
        # Odd runs are AUX input
        if self.training and self.num_batches_tracked % 2 != 0:
            return self.bn_aux(input)
      
        return self.bn(input)


bn = nn.BatchNorm2d(3)

with torch.no_grad():
    bn.weight.add_(10)
    bn.bias.add_(10)
    bn.running_mean.add_(10)
    bn.running_var.add_(10)

my_bn = MyBatchNorm2d(copy.deepcopy(bn))

x = torch.randn(2, 3, 4, 4)
out1 = bn(x)
out2 = my_bn(x)
print(out1)
print(out2)
</code></pre>
<p>How are you comparing the results using the test error?</p>

ialhashim: <p>I think it is working now. I tested it by running one branch on the original image and the other on the negative (-image) and the stats are now negative of each other after many iterations.</p>

Sara_Rojas_Martinez: <p>Hello,<br>
<a class="mention" href="/u/ialhashim">@ialhashim</a> Could you implement Advprop? I’m also having issues. <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=9" title=":frowning:" class="emoji" alt=":frowning:"></p>

ialhashim: <p><a class="mention" href="/u/sara_rojas_martinez">@Sara_Rojas_Martinez</a> I ended up with this code that comes with helper debug functions. Not sure if its correct as my experiments ended up learning two distribution equally well and not have a compounded advantage as shown in the paper.</p>
<pre><code class="lang-auto">class MyBatchNorm2d(nn.Module):
    def __init__(self, init_bn, isAuxBN):
        super(MyBatchNorm2d, self).__init__()
        self.bn = copy.deepcopy(init_bn)
        
        # Aux BN
        self.bn_aux = copy.deepcopy(init_bn)
        
        self.isAuxBN = isAuxBN
        self.isForwardAuxBN = False
        
        self.register_buffer('num_batches_tracked', torch.tensor(0, dtype=torch.long))
        
    def forward(self, input):            
        if self.training:
            self.num_batches_tracked = self.num_batches_tracked + 1            
        
        # Odd runs are AUX input
        if self.isAuxBN:
            if self.training and self.num_batches_tracked % 2 != 0:
                return self.bn_aux(input)
            
            if self.isForwardAuxBN:
                return self.bn_aux(input)
        
        return self.bn(input)
    
    def run_stats(self):
        return '\n\nParent {} \t bn {} \t bn_aux {}\n\n'.format(self.num_batches_tracked, self.bn.num_batches_tracked, self.bn_aux.num_batches_tracked)
        
    def bn_stats(self):
        return 'BN mean\n {}, \n\n'.format(self.bn.running_mean)
    
    def bn_aux_stats(self):
        return 'BN-AUX mean\n {}, \n'.format(self.bn_aux.running_mean)
    
    def __repr__(self):
        parent_string = super().__repr__()
        txt = parent_string + self.run_stats()
        txt = txt + self.bn_stats()
        txt = txt + self.bn_aux_stats() 
        return txt + '(is forward using AUX {})'.format(self.isForwardAuxBN)
</code></pre>

Petrea_Tancau1: <p>hi <a class="mention" href="/u/ialhashim">@ialhashim</a>, how did you generate the adversarial examples?<br>
Have you managed to apply PGD?</p>