public release

mit-han-lab · Dec 12, 2019 · d5dd20c · d5dd20c
commit d5dd20c
Show file tree

Hide file tree

Showing 10 changed files with 448 additions and 0 deletions.
diff --git a/.gitignore b/.gitignore
@@ -0,0 +1,104 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# pyenv
+.python-version
+
+# celery beat schedule file
+celerybeat-schedule
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
diff --git a/LICENSE b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2019 MIT HAN Lab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/README.md b/README.md
@@ -0,0 +1,86 @@
+# Deep Leakage From Gradients [[arXiv]](https://arxiv.org/abs/1906.08935) [[Webside]](https://dlg.mit.edu) 
+
+```
+@inproceedings{zhu19deep,
+  title={Deep Leakage from Gradients},
+  author={Zhu, Ligeng and Liu, Zhijian and Han, Song},
+  booktitle={Advances in Neural Information Processing Systems},
+  year={2019}
+}
+```
+
+Gradients exchaging is popular used in modern multi-node learning systems. People used to believe numerical gradients are safe to share. But we show that it is actually possible to obtain the training data from shared gradients and the leakage is pixel-wise accurate for images and token-wise matching for texts.
+
+<p align="center">
+    <img src="assets/nips-dlg.jpg" width="80%" />
+</p>
+
+<p align="center">
+    <img src="assets/demo-crop.gif" width="80%" />
+</p>
+
+## Overview
+
+We release the PyTorch code of [Deep Leakage from Gradients](https://arxiv.org/abs/1906.08935).
+
+<p align="center">
+    <img src="assets/method.jpg" width="80%" />
+</p>
+
+The core algorithm is to *match the gradients* between *dummy data* and *real data*. It can be implemented in **less than 20 lines**!
+
+
+```python
+def deep_leakage_from_gradients(model, origin_grad): 
+  dummy_data = torch.randn(origin_data.size())
+  dummy_label =  torch.randn(dummy_label.size())
+  optimizer = torch.optim.LBFGS([dummy_data, dummy_label] )
+
+  for iters in range(300):
+    def closure():
+      optimizer.zero_grad()
+      dummy_pred = model(dummy_data) 
+      dummy_loss = criterion(dummy_pred, dummy_label) 
+      dummy_grad = grad(dummy_loss, model.parameters(), create_graph=True)
+
+      grad_diff = sum(((dummy_grad - origin_grad) ** 2).sum() \
+        for dummy_g, origin_g in zip(dummy_grad, origin_grad))
+
+      grad_diff.backward()
+      return grad_diff
+
+    optimizer.step(closure)
+
+  return  dummy_data, dummy_label
+```
+
+
+# Prerequisites
+
+To run the code, following libraies are required
+
+* Python >= 3.6
+* PyTorch >= 1.0
+* torchvision >= 0.4
+
+# Code
+
+* If you do not have GPU mahcines, We provide [Google Colab](https://colab.research.google.com/gist/Lyken17/91b81526a8245a028d4f85ccc9191884/deep-leakage-from-gradients.ipynb) to quickly reproduce our results. 
+* If you have GPU servers and would like to run your locally, `python main.py` provides the same functionality.
+
+# Deep Leakage on Batched Images
+
+<p align="center">
+    <img src="assets/out.gif" width="80%" />
+</p>
+
+# Deep Leakage on Lanuage Model
+
+<p align="center">
+    <img src="assets/nlp_results.png" width="80%" />
+</p>
+
+
+# License
+
+This repository is released under the MIT license. See LICENSE for additional details.
diff --git a/assets/demo-crop.gif b/assets/demo-crop.gif
diff --git a/assets/method.jpg b/assets/method.jpg
diff --git a/assets/nips-dlg.jpg b/assets/nips-dlg.jpg
diff --git a/assets/nlp_results.png b/assets/nlp_results.png
diff --git a/assets/out.gif b/assets/out.gif
diff --git a/main.py b/main.py
@@ -0,0 +1,85 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+from pprint import pprint
+
+from PIL import Image
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import grad
+import torchvision
+from torchvision import models, datasets, transforms
+print(torch.__version__, torchvision.__version__)
+
+device = "cpu"
+if torch.cuda.is_available():
+    device = "cuda"
+print("Running on %s" % device)
+
+dst = datasets.CIFAR100("~/.torch", download=True)
+tp = transforms.ToTensor()
+tt = transforms.ToPILImage()
+
+img_index = 25
+gt_data = tp(dst[img_index][0]).to(device)
+gt_data = gt_data.view(1, *gt_data.size())
+gt_label = torch.Tensor([dst[img_index][1]]).long().to(device)
+gt_label = gt_label.view(1, )
+
+plt.imshow(tt(gt_data[0].cpu()))
+
+from model.vision import LeNet
+net = LeNet().to(device)
+
+net.apply(weights_init)
+criterion = nn.CrossEntropyLoss()
+
+# compute original gradient 
+out = net(gt_data)
+y = criterion(out, gt_label)
+dy_dx = torch.autograd.grad(y, net.parameters())
+
+original_dy_dx = list((_.detach().clone() for _ in dy_dx))
+original_dy_dx_fp16 = list((_.detach().clone().half().float() for _ in dy_dx))
+
+# generate dummy data and label
+dummy_data = torch.randn(gt_data.size()).to(device).requires_grad_(True)
+plt.imshow(tt(dummy_data[0].cpu()))
+
+optimizer = torch.optim.LBFGS([dummy_data, ] )
+
+condition = True
+prev_loss = 0
+counter = 0 
+
+history = []
+for iters in range(300):
+    def closure():
+        optimizer.zero_grad()
+
+        pred = net(dummy_data) 
+        dummy_loss = criterion(pred, gt_label) 
+        dummy_dy_dx = torch.autograd.grad(dummy_loss, net.parameters(), create_graph=True)
+
+        grad_diff = 0
+        for gx, gy in zip(dummy_dy_dx, original_dy_dx): 
+            grad_diff += ((gx - gy) ** 2).sum()
+        grad_diff.backward()
+
+        return grad_diff
+
+    optimizer.step(closure)
+    if iters % 10 == 0: 
+        current_loss = closure()
+        print(iters, "%.4f" % current_loss.item())
+        history.append(tt(dummy_data[0].cpu()))
+
+plt.figure(figsize=(12, 8))
+for i in range(30):
+    plt.subplot(3, 10, i + 1)
+    plt.imshow(history[i])
+    plt.title("iter=%d" % (i * 10))
+    plt.axis('off')
+