Calculating gradients when training=False in the callback (Solved)

I’m trying to implement adversarial test with callback. I have a little trouble getting the gradient.
Although I set requires_grad=True, I cannot get the gradients with respect to the input when training=False.

class FGSM(LearnerCallback):
    def __init__(self):
        super().__init__(learn)
        self.model = learn.model
        
    def on_batch_begin(self, **kwargs):
        X = kwargs['last_input']
        y = kwargs['last_target']
        if not kwargs['train']:
            X = self.fgsm(X, y)
            return {'last_input': X}
    
    def fgsm(self, X, y, epsilon=8/255):
        X = X.clone().detach()
        X = denormalize(X, torch.tensor(cifar_stats[0]), torch.tensor(cifar_stats[1]))
        # after denormalizing, X is not in GPU anymore, so move to cuda again
        X = to_device(X, device)
        X = X.requires_grad_(True)
        out = self.model(X) # out doesn't get any grad_fn although X's requires_grad=True
        losses = nn.CrossEntropyLoss(reduction='none')(out, y)
        loss = torch.mean(losses)
        # the grad cannot run because loss doesn't have grad_fn
        # if i explicitly set loss.requires_grad_(True), the gradients become None
        grad, _ = torch.autograd.grad(loss, [X])
        X = X + torch.sign(grad) * epsilon
        X = torch.clamp(X, 0, 1)
        X = normalize(X, torch.tensor(cifar_stats[0]), torch.tensor(cifar_stats[1]))
        return X

May I know some tips on how to get gradients with respect to the input in callbacks? Thank you very much.

1 Like

I am not quite sure what you are trying to do here, so my input may very well be useless. Let me know if it is, and I will remove it.

To access the gradient, you grab the grad or grad.data attribute of a pytorch Tensor. So if i wanted the gradient of one of the model parameters, i would use the parameters iterator. In your case: list(self.model.parameters())[0].grad would give you the first parameter’s gradient.

I think your problem is: you have to run the backward() method on a tensor to populate the gradient attribute. Normally this is done on the loss variable, hence the line in the training loop loss.backward(). This line results in any tensors with requires_grad=True, involved in the computation of the loss having their grad attribute populated. It does not look like you are running backward() on any tensors, nor are you waiting for the training loop to do this for you.

In fastai, without calculating loss yourself, you would need to wait until after loss.backward() is called in the training loop in order to get a gradient. Your callback would need to be in between the loss.backward() and opt.step(); opt.zero_grad() calls. Otherwise, the training loop will do an optimizer step, and then zero all the gradients.

Also, why are you trying to calculate loss in a callback? The main training loop supports custom loss functions on its own?

I hope this helps. :smile:

2 Likes

@exynos7 Thank you very much for your reply.
What I am trying to do is I want to add adversarial noise to the validation batch.
So, I need the gradient with respect to the input (from the code above that is X), (not the parameters from the model).

The problem is: the output of self.model(input) is not getting grad_fn=<AddmmBackward> when training=False.
Therefore, I cannot run loss.backward() or autograd.grad because loss does not get grad_fn=<MeanBackward0> too. I have tried to do it during training (training=True), everything works fine as I want. But I need to do it when training=False.

The overall idea is I want to calculate the gradients with respect to the input during validation. With that gradient sign, I want to add noise to the input batch.

I have solved it. Thank you anyway.
The problem was torch.no_grad() is doing the trick. In fastai, torch.no_grad() is set since the beginning in the validation. I put torch.enable_grad() to do what I need to do. Everything works perfectly. I love fastai. Thank you all. To be honest, because of callbacks, I started using fastai. It’s amazing.

@pmaung Is your FGSM example in fast.ai available for sharing? I have seen FGSM tutorials in Tensorflow and Pytorch and it would be great to see who you implemented it using callbacks.

Thanks,
Jeff

@jeffchen72 Sure. Here it is.

class FGSM(LearnerCallback):
    def __init__(self):
        super().__init__(learn)
        self.model = learn.model
        
    def on_batch_begin(self, **kwargs):
        X = kwargs['last_input'].clone()
        y = kwargs['last_target']
        if not kwargs['train']:       
            X = self.fgsm(X, y)
            return {'last_input': X}
    
    def fgsm(self, X, y, epsilon=64/255):
        x = X.clone().detach()
        mean = torch.tensor(imagenet_stats[0])
        std = torch.tensor(imagenet_stats[1])
        x = denormalize(x, mean, std)
        x = to_device(x, device)
        x = x.requires_grad_(True)
        with torch.enable_grad():
            out = self.model(x)
            losses = nn.CrossEntropyLoss(reduction='none')(out, y)
            loss = torch.mean(losses)
            grad, = torch.autograd.grad(loss, [x])
        x = x + torch.sign(grad) * epsilon
        x = torch.clamp(x, 0, 1)
        x = normalize(x, mean.cuda(), std.cuda())
        return x

To call it, on your trained model, for example

learn.validate(data.valid_dl, callbacks=[FGSM()])

It works with me. Please let me know if it is not clear.

Thanks @pmaung for sharing your code. I am still learning fast.ai, so please forgive my novice questions. Once we generate the perturbation using FGSM to build an adversarial example, how do we access it to attack the trained model?

I would like to replicate this Pytorch fgsm example using fast.ai as a way to learn.

https://pytorch.org/tutorials/beginner/fgsm_tutorial.html

Here’s my code with your FGSM callback incorporated is below. Thanks in advance for your help.

from fastai.vision import *

# FGSM throws error if I don't set GPU=1 on my two GPU system
device = torch.device(0)

path = untar_data(URLs.MNIST_SAMPLE)
data = ImageDataBunch.from_folder(path)

# Callback for generating FGSM perturbation
class FGSM(LearnerCallback):
    def __init__(self):
        super().__init__(learn)
        self.model = learn.model
        
    def on_batch_begin(self, **kwargs):
        X = kwargs['last_input'].clone()
        y = kwargs['last_target']
        if not kwargs['train']:
            X = self.fgsm(X, y)
            return {'last_input': X}
    
    def fgsm(self, X, y, epsilon=8/255):
        x = X.clone().detach()
        # Change depending on your dataset: imagenet_stats, cifar_stats, mnist_stats...
        mean = torch.tensor(mnist_stats[0]) 
        std = torch.tensor(mnist_stats[1])
        x = denormalize(x, mean, std)
        x = to_device(x, device)
        x = x.requires_grad_(True)
        with torch.enable_grad():
            out = self.model(x)
            losses = nn.CrossEntropyLoss(reduction='none')(out, y)
            loss = torch.mean(losses)
            grad, = torch.autograd.grad(loss, [x])
        x = x + torch.sign(grad) * epsilon
        x = torch.clamp(x, 0, 1)
        x = normalize(x, mean.cuda(), std.cuda())
        return x
    
# After training model, call
#     learn.validate(data.valid_dl, callbacks=[FGSM()])

# Create a learner
learn = cnn_learner(data, models.resnet18, metrics=accuracy)

# Train model
learn.fit_one_cycle(1)

# Interpret the model
# interp = ClassificationInterpretation.from_learner(learn)
# losses,idxs = interp.top_losses()
# interp.plot_top_losses(9, figsize=(15,11))
# interp.plot_confusion_matrix(figsize=(12,12), dpi=60)
learn.show_results()

learn.validate(data.valid_dl)

pred = learn.predict(learn.data.valid_ds[0][0])
pred

learn.data.train_ds[0][0]

# Generate FGSM perturbation
learn.validate(data.valid_dl, callbacks=[FGSM()])

# How do you access the perturbed adversarial data to attack the trained model?
#
# adv_data = FSGM.X
# pred = learn.predict(adv_data[0][0])
# pred

@jeffchen72 I’m not sure about multiple GPUs with fastai. I have never tried. Do you want to see adversarial examples? I didn’t quite get your questions. When you pass a callback, you are basically attacking. learn.validate will use adversarial examples to test if you give it FGSM callback. The tutorial and the callback is doing exactly the same thing. You can calculate the gradients as you like whether with loss.backward() or torch.autograd.grad. I stole the that line from https://github.com/MadryLab/robustness_lib/blob/master/robustness/attacker.py.

I make a simple notebook to show you an example.
https://github.com/fugokidi/study/blob/master/fgsm.ipynb

I am also a beginner. Please let me know if you find a mistake.

@pmaung Thank you so much for sharing your notebook. There are some things that I don’t quite understand. It seems when you call

learn.validate(data.valid_dl, callbacks=[FGSM()])

it generates calculated loss and metrics, but the data in the dl is not modified with the fgsm attack. That is, if you run

x, y = learn.get_preds(data.valid_dl) # get predictions on the validation set
x[0]

You get the same results before and after running

learn.validate(data.valid_dl, callbacks=[FGSM()])

However, when you run

x_adv = FGSM().fgsm(x, y, epsilon=0.10)

x_adv contains the adversarial attack and showing the image clears shows the modifications.

I am would be interested in comparing the results learn.predict() before and after the fgsm attack.

However,

learn.predict(x_adv)

results in an error.

AttributeError: 'Tensor' object has no attribute 'apply_tfms'

How do I put x_adv into the correct datatype to run inference on the adversarial example?

Thanks,
Jeff

@jeffchen72 Well, seems like you don’t understand me and I don’t understand you, lol…just kidding.

  1. Let me explain a bit about callback. Callbacks are used to intercept the training loop and do whatever you like. They never modify anything to your dataset. When I pass the callback, it modifies the the batch on the fly.
    I was showing you an example if you want to see adversarial examples by passing a batch to FGSM().fgsm(x, y), x is a batch.

  2. You cannot call learn.predict(x_adv) because predict() is expecting a specific piece of data (image data), not a tensor object. That’s why you got that error. If you want to call predict, run the followings.
    learn.predict(Image(x[0]) for clean image test
    learn.predict(Image(x_adv[0]) for adversarial example

Hope it helps.

@pmaung Thanks for your patience. :smiley:

I understand it now. I tried it out and it is working. These are some errors that I encountered, but I also found workarounds. Not sure if they are bugs or my errors.

https://github.com/jc7k/study/blob/master/cifar10-fgsm-fastai.ipynb

This is a good learning experience for me. Thanks for your help.