Fast Gradient Sign Method (FGSM) Example for MNIST Using PyTorch

The fast gradient sign method (FGSM) is a technique that takes an input item (usually an image) and generates an evil near-copy of the item that has been deliberately constructed to produce an incorrect prediction.

In the screenshot below, the demo trains a model to classify the MNIST handwritten digits dataset. The model scores 98.00% accuracy on the 100-item test dataset. Then the demo takes the 100 test images and creates 100 evil copies of the test items. The model scores only 29% accuracy on those evil images even though they look very similar to the source test images.


The top image is a ‘2’ from the test data. The bottom image is an evil ‘2’ created from the test ‘2’ using FGSM. The trained model thinks the evil image is a ‘6’.

In the screenshot, the top ‘2’ is the test image [1] (the second image). The bottom image was generated using FGSM. It looks like a ‘2’ to the human eye, but the trained model incorrectly classifies the image as a ‘6’.

Conceptally, FGSM is quite simple, but getting the demo up and running took me a few hours. There were many tricky details. During training, the model weights are updated roughly like this:

compute all weight gradients
for-each weight
  new wt = old wt - lrn_rate * gradient
end-for

For training, if the model has 10,000 weights, there are 10,000 corresponding gradients. The subtraction adjusts the weights to make the model give a better prediction.

In FGSM, you add gradients for each input pixel, so for the 28×28 = 784 input values there are 784 input gradients. Creating the evil copy looks roughly like:

compute all input pixel gradients
for-each input pixel
  new pixel = old pixel + epsilon * gradient
end-for

The addition adjusts the input pixels to make the model give a worse prediction. Large values of epsilon make the model produce more incorrect predictions, but the evil near-copy looks less like its source.

Fascinating stuff.



A common theme in science fiction movies is alien copies of humans. Left: In “The Faculty” (1998), alien parasites take over a high school. The aliens are defeated in the end. Center: In “It Came From Outer Space” (1953), aliens crash in the desert and transform themselves to look like townspeople. But it turns out the aliens are friendly and just want to repair their spaceship. They do and there’s a happy ending. Right: In “Species” (1995), scientists splice alien DNA with human DNA. The result was pleasing to the eye but deadly to the men who were a target of the alien’s mating urge.


Demo code. Replace “lt”, “gt”, “lte”, “gte” with Boolean operator symbols.

# mnist_fgsm.py
# PyTorch 1.10.0-CPU Anaconda3-2020.02  Python 3.7.6
# Windows 10/11

# generate adversarial data using the fast gradient
# sign method (FGSM)

# reads MNIST data from text file rather than using
# built-in black box Dataset from torchvision
# see jamesmccaffrey.wordpress.com/2022/01/21/ + 
#       working-with-mnist-data/
# see jamesmccaffrey.wordpress.com/2021/03/15/ +
#       converting-mnist-binary-files-to-text-files/

import numpy as np
import matplotlib.pyplot as plt
import torch as T
import matplotlib.pyplot as plt

device = T.device('cpu')

# -----------------------------------------------------------

class MNIST_Dataset(T.utils.data.Dataset):
  # 784 tab-delim pixel values (0-255) then label (0-9)
  def __init__(self, src_file):
    all_xy = np.loadtxt(src_file, usecols=range(785),
      delimiter="\t", comments="#", dtype=np.float32)

    tmp_x = all_xy[:, 0:784]  # all rows, cols [0,783]
    tmp_x /= 255.0
    tmp_x = tmp_x.reshape(-1, 1, 28, 28)  # bs, chnls, 28x28
    tmp_y = all_xy[:, 784]    # 1-D required

    self.x_data = \
      T.tensor(tmp_x, dtype=T.float32).to(device)
    self.y_data = \
      T.tensor(tmp_y, dtype=T.int64).to(device) 

  def __len__(self):
    return len(self.x_data)

  def __getitem__(self, idx):
    lbl = self.y_data[idx] 
    pixels = self.x_data[idx] 
    return (pixels, lbl)

# -----------------------------------------------------------

class Net(T.nn.Module):
  def __init__(self):
    super(Net, self).__init__()  # pre Python 3.3 syntax

    self.conv1 = T.nn.Conv2d(1, 32, 5)  # chnl-in, out, krnl
    self.conv2 = T.nn.Conv2d(32, 64, 5)
    self.fc1 = T.nn.Linear(1024, 512)   # [64*4*4, x]
    self.fc2 = T.nn.Linear(512, 256)
    self.fc3 = T.nn.Linear(256, 10)     # 10 classes
    self.pool1 = T.nn.MaxPool2d(2, 2)   # kernel, stride
    self.pool2 = T.nn.MaxPool2d(2, 2)
    self.drop1 = T.nn.Dropout(0.25)
    self.drop2 = T.nn.Dropout(0.50)
    # default weight and bias initialization
  
  def forward(self, x):
    # convolution phase         # x is [bs, 1, 28, 28]
    z = T.relu(self.conv1(x))   # Size([bs, 32, 24, 24])
    z = self.pool1(z)           # Size([bs, 32, 12, 12])
    z = self.drop1(z)
    z = T.relu(self.conv2(z))   # Size([bs, 64, 8, 8])
    z = self.pool2(z)           # Size([bs, 64, 4, 4])
   
    # neural network phase
    z = z.reshape(-1, 1024)     # Size([bs, 1024])
    z = T.relu(self.fc1(z))     # Size([bs, 512])
    z = self.drop2(z)
    z = T.relu(self.fc2(z))     # Size([bs, 256])
    z = self.fc3(z)             # Size([bs, 10])
    return z

# -----------------------------------------------------------

def accuracy(model, ds):
  ldr = T.utils.data.DataLoader(ds,
    batch_size=len(ds), shuffle=False)
  n_correct = 0
  for data in ldr:
    (pixels, labels) = data
    with T.no_grad():
      oupts = model(pixels)
    (_, predicteds) = T.max(oupts, 1)
    n_correct += (predicteds == labels).sum().item()

  acc = (n_correct * 1.0) / len(ds)
  return acc

# -----------------------------------------------------------

def main():
  # 0. setup
  print("\nBegin MNIST with PyTorch FGSM demo ")
  np.random.seed(1)
  T.manual_seed(1)

  # 1. create Dataset
  print("\nCreating 1000-item train Dataset from text file ")
  train_file = ".\\Data\\mnist_train_1000.txt"
  train_ds = MNIST_Dataset(train_file)

  print("Creating 100-item test Dataset from text file ")
  test_file = ".\\Data\\mnist_test_100.txt"
  test_ds = MNIST_Dataset(test_file)

  bat_size = 20
  train_ldr = T.utils.data.DataLoader(train_ds,
    batch_size=bat_size, shuffle=True)

  # 2. create network
  print("\nCreating CNN network with 2 conv and 3 linear ")
  net = Net().to(device)

# -----------------------------------------------------------

  # 3. train model
  max_epochs = 25  # 100 gives better results
  ep_log_interval = 5
  lrn_rate = 0.05
  
  loss_func = T.nn.CrossEntropyLoss()  # does log-softmax()
  optimizer = T.optim.SGD(net.parameters(), lr=lrn_rate)
    
  print("\nbat_size = %3d " % bat_size)
  print("loss = " + str(loss_func))
  print("optimizer = SGD")
  print("lrn_rate = %0.3f " % lrn_rate)
  print("max_epochs = %3d " % max_epochs)

  print("\nStarting training")
  net.train()  # set mode
  for epoch in range(0, max_epochs):
    ep_loss = 0  # for one full epoch
    for (batch_idx, batch) in enumerate(train_ldr):
      (X, y) = batch  # X = pixels, y = target labels
      optimizer.zero_grad()
      oupt = net(X)
      loss_val = loss_func(oupt, y)  # a tensor
      ep_loss += loss_val.item()  # accumulate
      loss_val.backward()  # compute grads
      optimizer.step()     # update weights
    if epoch % ep_log_interval == 0:
      print("epoch = %4d   |  loss = %9.4f" % (epoch, ep_loss))
  print("Done ") 

# -----------------------------------------------------------

  # 4. evaluate model accuracy
  print("\nComputing model accuracy")
  net.eval()
  acc_train = accuracy(net, train_ds)  # all at once
  print("Accuracy on training data = %0.4f" % acc_train)
  
  net.eval()
  acc_test = accuracy(net, test_ds)  # all at once
  print("Accuracy on test data = %0.4f" % acc_test)

# -----------------------------------------------------------

  # 5. use model to make prediction: N/A
  
# -----------------------------------------------------------

  # 6. save model
  # print("\nSaving trained model state")
  # fn = ".\\Models\\mnist_model.pt"
  # T.save(net.state_dict(), fn)  

# -----------------------------------------------------------

  # 7. create inputs designed to trick model
  epsilon = 0.20
  print("\nCreating evil images from test with epsilon = \
    %0.2f " % epsilon)
  evil_images_lst = []
  n_correct = 0; n_wrong = 0

  test_ldr = T.utils.data.DataLoader(test_ds,
    batch_size=1, shuffle=False)
  loss_func = T.nn.CrossEntropyLoss()  # does log-softmax()

  for (batch_idx, batch) in enumerate(test_ldr):
    (X, y) = batch  # X = pixels, y = target label
    X.requires_grad = True
    oupt = net(X)
    loss_val = loss_func(oupt, y)
    net.zero_grad()  # zap all gradients
    loss_val.backward()  # compute gradients

    sgn = X.grad.data.sign()
    mutated = X + epsilon * sgn
    mutated = T.clamp(mutated, 0.0, 1.0)

    with T.no_grad():
      pred = net(mutated)  # 10 logits
    pred_class = T.argmax(pred[0])
    if pred_class.item() == y.item():
      n_correct += 1
    else:
      n_wrong += 1

    mutated = mutated.detach().numpy()
    evil_images_lst.append(mutated)
    
  # print(n_correct)
  # print(n_wrong)
  adver_acc = (n_correct * 1.0) / (n_correct + n_wrong)
  print("\nModel acc on evil images = %0.4f " % adver_acc)
  
  # show first test image and corresponding mutation
  pixels = test_ds[0][0].reshape(28,28)
  plt.imshow(pixels, cmap=plt.get_cmap('gray'))
  plt.show() 

  pixels = evil_images_lst[0].reshape(28,28)
  plt.imshow(pixels, cmap=plt.get_cmap('gray'))
  plt.show() 
 
  print("\nEnd MNIST PyTorch FGSM demo ")

if __name__ == "__main__":
  main()
This entry was posted in PyTorch. Bookmark the permalink.