resnet_fgsm.py

import torch
from torchvision.models.resnet import resnet18
from torchvision.utils import save_image
from base64 import b64encode as b64
from Embedding import key
from io import BytesIO

device='cuda'

model = resnet18(pretrained=True).to(device)

#get embedding rather than logits from final layer
model.fc = torch.nn.Identity()

key = torch.tensor(key).to(device)
tensored = torch.rand(1, 3, 224, 224).to(device)

# https://pytorch.org/tutorials/beginner/fgsm_tutorial.html#fgsm-attack
while True:
    tensored.requires_grad_()
    embedding = model(tensored)[0]
    diff = ((embedding - key)**2).mean()
    if diff.item() < 1e-4:
        b = BytesIO()
        save_image(tensored[0], b, format="png")
        print(b64(b.getvalue()).decode())
        break
    
    diff.backward()
    tensored = tensored.detach() - tensored.grad * 64