forked from jonasrauber/eagerpy
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
added docs for converting, generic functions, and autodiff
- Loading branch information
Jonas Rauber
committed
Aug 14, 2020
1 parent
17fcc1e
commit a222f39
Showing
4 changed files
with
152 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,30 @@ | ||
--- | ||
title: Automatic Differentiation | ||
|
||
--- | ||
|
||
# Automatic Differentiation in EagerPy | ||
|
||
EagerPy uses a functional approach to automatic differentiation. You first define a function that will then be differentiated with respect to its inputs. This function is then passed to `ep.value_and_grad` to evaluate both the function and its gradient. More generally, you can also use `ep.value_aux_and_grad` if your function has additional auxiliary outputs and `ep.value_and_grad_fn` if you want the gradient function without immediately evaluating it at some point `x`. | ||
|
||
Using `ep.value_and_grad` for automatic differentiation in EagerPy: | ||
|
||
```python | ||
import torch | ||
x = torch.tensor([1., 2., 3.]) | ||
|
||
# The following code works for any framework, not just Pytorch! | ||
|
||
import eagerpy as ep | ||
x = ep.astensor(x) | ||
|
||
def loss_fn(x): | ||
# this function takes and returns an EagerPy tensor | ||
return x.square().sum() | ||
|
||
print(loss_fn(x)) | ||
# PyTorchTensor(tensor(14.)) | ||
|
||
print(ep.value_and_grad(loss_fn, x)) | ||
# (PyTorchTensor(tensor(14.)), PyTorchTensor(tensor([2., 4., 6.]))) | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,65 @@ | ||
--- | ||
title: Converting | ||
|
||
--- | ||
|
||
# Converting Between EagerPy and Native Tensors | ||
|
||
A native tensor could be a PyTorch GPU or CPU tensor, a TensorFlow tensor, a JAX array, or a NumPy array. | ||
|
||
**A native PyTorch tensor:** | ||
```python | ||
import torch | ||
x = torch.tensor([1., 2., 3., 4., 5., 6.]) | ||
``` | ||
|
||
**A native TensorFlow tensor:** | ||
```python | ||
import tensorflow as tf | ||
x = tf.constant([1., 2., 3., 4., 5., 6.]) | ||
``` | ||
|
||
**A native JAX array:** | ||
```python | ||
import jax.numpy as np | ||
x = np.array([1., 2., 3., 4., 5., 6.]) | ||
``` | ||
|
||
**A native NumPy array:** | ||
```python | ||
import numpy as np | ||
x = np.array([1., 2., 3., 4., 5., 6.]) | ||
``` | ||
|
||
No matter which native tensor you have, it can always be turned into the appropriate EagerPy tensor using `ep.astensor`. This will automatically wrap the native tensor with the correct EagerPy tensor class. The original native tensor can always be accessed using the `.raw` attribute. | ||
|
||
```python | ||
# x should be a native tensor (see above) | ||
# for example: | ||
import torch | ||
x = torch.tensor([1., 2., 3., 4., 5., 6.]) | ||
|
||
# Any native tensor can easily be turned into an EagerPy tensor | ||
import eagerpy as ep | ||
x = ep.astensor(x) | ||
|
||
# Now we can perform any EagerPy operation | ||
x = x.square() | ||
|
||
# And convert the EagerPy tensor back into a native tensor | ||
x = x.raw | ||
# x will now again be a native tensor (e.g. a PyTorch tensor) | ||
``` | ||
|
||
Especially in functions, it is common to convert all inputs to EagerPy tensors. This could be done using individual calls to `ep.astensor`, but using `ep.astensors` this can be written even more compactly. | ||
|
||
```python | ||
# x, y should be a native tensors (see above) | ||
# for example: | ||
import torch | ||
x = torch.tensor([1., 2., 3.]) | ||
y = torch.tensor([4., 5., 6.]) | ||
|
||
import eagerpy as ep | ||
x, y = ep.astensors(x, y) # works for any number of inputs | ||
``` |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
--- | ||
title: Generic Functions | ||
|
||
--- | ||
|
||
# Implementing Generic Framework-Agnostic Functions | ||
|
||
Using the conversion functions shown in \appendixref{sec:converting}, we can already define a simple framework-agnostic function. This function can be called with a native tensor from any framework and it will return the norm of that tensor, again as a native tensor from that framework (\autoref{alg:calling_norm_pytorch}, \autoref{alg:calling_norm_tensorflow}). | ||
|
||
```pyhon | ||
import eagerpy as ep | ||
def norm(x): | ||
x = ep.astensor(x) | ||
result = x.square().sum().sqrt() | ||
return result.raw | ||
``` | ||
|
||
Calling the `norm` function using a PyTorch tensor: | ||
```pyhon | ||
import torch | ||
norm(torch.tensor([1., 2., 3.])) | ||
# tensor(3.7417) | ||
``` | ||
|
||
Calling the `norm` function using a TensorFlow tensor: | ||
```python | ||
import tensorflow as tf | ||
norm(tf.constant([1., 2., 3.])) | ||
# <tf.Tensor: shape=(), dtype=float32, numpy=3.7416575> | ||
``` | ||
|
||
If we would call the above `norm` function with an EagerPy tensor, the `ep.astensor` call would simply return its input. The `result.raw` call in the last line would however still extract the underlying native tensor. Often it is preferably to implement a generic function that not only transparently handles any native tensor but also EagerPy tensors, that is the return type should always match the input type. This is particularly useful in libraries like Foolbox that allow users to work with EagerPy and native tensors. To achieve that, EagerPy comes with two derivatives of the above conversion functions: `ep.astensor_` and `ep.astensors_`. Unlike their counterparts without an underscore, they return an additional inversion function that restores the input type. If the input to `astensor_` is a native tensor, `restore_type` will be identical to `.raw`, but if the original input was an EagerPy tensor, `restore_type` will not call `.raw`. With that, we can write generic framework-agnostic functions that work transparently for any input. | ||
|
||
|
||
An improved framework-agnostic `norm` function: | ||
```python | ||
import eagerpy as ep | ||
|
||
def norm(x): | ||
x, restore_type = ep.astensor_(x) | ||
result = x.square().sum().sqrt() | ||
return restore_type(result) | ||
``` | ||
|
||
Converting and restoring multiple inputs using `ep.astensors_`: | ||
```python | ||
import eagerpy as ep | ||
|
||
def example(x, y, z): | ||
(x, y, z), restore_type = ep.astensors_(x, y, z) | ||
result = (x + y) * z | ||
return restore_type(result) | ||
``` |