Spaces:
Build error
Build error
| import json | |
| import torch | |
| import torchvision.transforms as transforms | |
| import torch.nn.functional as F | |
| import requests | |
| from .base_env import device | |
| def biased_fgsm(model, image): | |
| # 计算梯度 | |
| image.requires_grad = True | |
| outputs = model(image.unsqueeze(0)) | |
| original_label = torch.argmax(outputs, 1) | |
| loss = F.cross_entropy(outputs, original_label) # IMPORTANT！！在此假设了原始标签是预测正确的，以此计算交叉熵损失后求梯度 | |
| loss.backward() | |
| # 添加对抗扰动 | |
| sngd = image.grad.sign() | |
| return sngd | |
| def predict(model, image): | |
| global labels | |
| with torch.no_grad(): | |
| prediction = F.softmax(model(image)[0], dim=0) | |
| confidences = {labels[i]: float(prediction[i]) for i in range(1000)} | |
| return confidences | |
| def show_attack(model_name, epsilon, PIL_image): | |
| # 创建一个新的预处理函数 | |
| preprocess = transforms.Compose([ | |
| transforms.Resize(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # ImageNet 的均值和标准差 | |
| ]) | |
| image = preprocess(PIL_image) | |
| if device == torch.device('cuda'): | |
| image = image.cuda() | |
| # 加载模型 | |
| if model_name == "AlexNet": | |
| model = torch.load('./models/PublicPretrained/alexnet.pth') | |
| elif model_name == "VGG16": | |
| model = torch.load('./models/PublicPretrained/vgg16.pth') | |
| elif model_name == "GoogLeNet": | |
| model = torch.load('./models/PublicPretrained/googlenet.pth') | |
| elif model_name == "ResNet18": | |
| model = torch.load('./models/PublicPretrained/resnet18.pth') | |
| elif model_name == "ResNet50": | |
| model = torch.load('./models/PublicPretrained/resnet50.pth') | |
| elif model_name == "Inception_V3": | |
| model = torch.load('./models/PublicPretrained/inception_v3.pth') | |
| elif model_name == "DenseNet161": | |
| model = torch.load('./models/PublicPretrained/densenet161.pth') | |
| elif model_name == "ResNeXt50_32x4d": | |
| model = torch.load('./models/PublicPretrained/resnext50_32x4d.pth') | |
| elif model_name == "MobileNet_V2": | |
| model = torch.load('./models/PublicPretrained/mobilenet_v2.pth') | |
| if device == torch.device('cuda'): | |
| model = model.cuda() | |
| model.eval() | |
| # Download human-readable labels for ImageNet. | |
| with open('./configs/imagenet_labels.json', 'r') as f: | |
| global labels | |
| labels = json.load(f) | |
| sngd = biased_fgsm(model, image) | |
| adv_image = image + epsilon * sngd | |
| adv_image = adv_image.detach() # 将adv_image从计算图中分离，不再计算梯度 | |
| unprocess = transforms.Compose([ | |
| transforms.Normalize(mean=[0, 0, 0], std=[1 / 0.229, 1 / 0.224, 1 / 0.225]), | |
| transforms.Normalize(mean=[-0.485, -0.456, -0.406], std=[1, 1, 1]), | |
| transforms.ToPILImage() | |
| ]) | |
| PIL_sngd = transforms.functional.to_pil_image(sngd * 0.5 + 0.5) | |
| PIL_adv_image = unprocess(adv_image) | |
| #print(str(device)) | |
| confidences = predict(model, image.unsqueeze(0).to(device)) | |
| sngd_confidences = predict(model, sngd.unsqueeze(0).to(device)) | |
| adv_confidences = predict(model, adv_image.unsqueeze(0).to(device)) | |
| pred, conf = max(confidences.items(), key=lambda x: x[1]) | |
| sngd_pred, sngd_conf = max(sngd_confidences.items(), key=lambda x: x[1]) | |
| adv_pred, adv_conf = max(adv_confidences.items(), key=lambda x: x[1]) | |
| top5_adv_confidences = dict(sorted(adv_confidences.items(), key=lambda x: x[1], reverse=True)[:10]) | |
| return [ | |
| top5_adv_confidences, | |
| PIL_image, pred, conf, | |
| PIL_sngd, sngd_pred, sngd_conf, | |
| PIL_adv_image, adv_pred, adv_conf | |
| ] | |
| if __name__ == "__main__": | |
| pass | |