ΦML Quickstart¶

Colab   •   🌐 ΦML   •   📖 Documentation   •   🔗 API   •   ▶ Videos   •   Examples

Installation¶

Install ΦML with pip on Python 3.6 and later:

In [1]:
%%capture
!pip install phiml

Install PyTorch, TensorFlow or Jax to enable machine learning capabilities and GPU execution. See the detailed installation instructions.

In [2]:
from phiml import math

Usage without ΦML's Tensors¶

You can call many functions on native tensors directly. ΦML will dispatch the call to the corresponding library and return the result as another native tensor.

In [3]:
math.sin(1.)
Out[3]:
0.841471
In [4]:
from jax import numpy as jnp
math.sin(jnp.asarray([1.]))

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Out[4]:
Array([0.84147096], dtype=float32)
In [5]:
import torch
math.sin(torch.tensor([1.]))
Out[5]:
tensor([0.8415])
In [6]:
import tensorflow as tf
math.sin(tf.constant([1.]))

2024-10-18 15:40:27.860015: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Out[6]:
<tf.Tensor: shape=(1,), dtype=float32, numpy=array([0.84147096], dtype=float32)>
In [7]:
import numpy as np
math.sin(np.asarray([1.]))
Out[7]:
array([0.841471], dtype=float32)

ΦML's Tensor¶

For more advanced operations, we recommend using ΦML's tensors. While ΦML includes a unified low-level API that behaves much like NumPy, using it correctly (so that the code is actually compatible with all libraries) is difficult. Instead, ΦML provides a higher-level API consisting of the Tensor class, the math functions and other odds and ends, that makes writing unified code easy. Tensors can be created by wrapping an existing backend-specific tensor or array:

In [8]:
torch_tensor = torch.tensor([1, 2, 3])
math.tensor(torch_tensor)
Out[8]:
(1, 2, 3) int64
In [9]:
math.wrap(torch_tensor)
Out[9]:
(1, 2, 3) int64

The difference between tensor and wrap is that wrap keeps the original data you pass in while tensor will convert the data to the default backend which can be set using math.use().

In [10]:
math.use('jax')
math.wrap(torch_tensor).default_backend
Out[10]:
torch
In [11]:
math.tensor(torch_tensor).default_backend
Out[11]:
jax

The last tensor call converted the PyTorch tensor to a Jax DeviceArray using a no-copy routine from dlpack under the hood.

Dimension Types¶

For tensors with more than one dimensions, you have to specify a name and type for each. Possible types are batch for parallelizing code, channel for listing features (color channels or x/y/z components) and spatial for equally-spaced sample points (width/height of an image, 1D time series, etc.). For an exhaustive list, see here

In [12]:
from phiml.math import batch, spatial, channel
torch_tensor = torch.tensor([[1, 2, 3], [4, 5, 6]])
math.wrap(torch_tensor, batch('dim1'), channel('dim2'))
Out[12]:
(1, 2, 3); (4, 5, 6) (dim1ᵇ=2, dim2ᶜ=3) int64

The superscript b and c denote the dimension type. When creating a new tensor from scratch, we also need to specify the size along each dimension:

In [13]:
math.random_uniform(batch(dim1=2), channel(dim2=3))
Out[13]:
(0.072, 0.020, 0.077); (0.879, 0.165, 0.102) (dim1ᵇ=2, dim2ᶜ=3)

When passing tensors to a neural network, the tensors are transposed to match the preferred dimension order (BHWC for TensorFlow/Jax, BCHW for PyTorch). For example, we can pass any number of batch and channel dimensions to an MLP.

In [14]:
from phiml import nn
mlp = nn.mlp(in_channels=6, out_channels=3, layers=[64, 64])
data = math.random_normal(batch(b1=4, b2=10), channel(c1=2, c2=3))
math.native_call(mlp, data)
Out[14]:
(b1ᵇ=4, b2ᵇ=10, vectorᶜ=3) -1.04e-04 ± 3.0e-01 (-1e+00...9e-01)

The network here is a standard fully-connected network module with two hidden layers of 64 neurons each. The native tensor that is passed to the network has shape (40, 6) as all batch dimensions are compressed into the first and all channel dimensions into the last dimension.

For a network acting on spatial data, we would add spatial dimensions.

In [15]:
net = nn.u_net(in_channels=6, out_channels=3, in_spatial=2)
data = math.random_normal(batch(b1=4, b2=10), channel(c1=2, c2=3), spatial(x=28, y=28))
math.native_call(mlp, data)
Out[15]:
(b1ᵇ=4, b2ᵇ=10, xˢ=28, yˢ=28, vectorᶜ=3) -0.004 ± 0.322 (-2e+00...2e+00)

In this example, we ran a 2D U-Net. For a 1D or 3D variant, we would pass in_spatial=1 or 3, respectively, and add the corresponding number of spatial dimensions to data.

Slicing¶

Slicing in ΦML is done by dimension names. Say we have a set of images:

In [16]:
images = math.random_uniform(batch(set=4), spatial(x=28, y=28), channel(channels=3))
images
Out[16]:
(setᵇ=4, xˢ=28, yˢ=28, channelsᶜ=3) 0.502 ± 0.288 (1e-04...1e+00)

The red, green and blue components are stored inside the channels dimension. Then to get just the red component of the last entry in the set, we can write

In [17]:
images.set[-1].channels[0]
Out[17]:
(xˢ=28, yˢ=28) 0.501 ± 0.291 (2e-03...1e+00)

Or we can slice using a dictionary

In [18]:
images[{'set': -1, 'channels': 0}]
Out[18]:
(xˢ=28, yˢ=28) 0.501 ± 0.291 (2e-03...1e+00)

Slicing the NumPy way, i.e. images[-1, :, :, 0] is not supported because the order of dimensions generally depends on which backend you use.

To make your code easier to read, you may name slices along dimensions as well. In the above example, we might name the red, green and blue channels explicitly:

In [19]:
images = math.random_uniform(batch(set=4), spatial(x=28, y=28), channel(channels='red,green,blue'))
images.set[-1].channels['red']
images[{'set': -1, 'channels': 'red'}]
Out[19]:
(xˢ=28, yˢ=28) 0.500 ± 0.276 (2e-03...1e+00)

To select multiple items by index, use the syntax tensor.<dim>[start:end:step] where start >= 0, end and step > 0 are integers.

In [20]:
images.x[1:3]
Out[20]:
(setᵇ=4, xˢ=2, yˢ=28, channelsᶜ=red,green,blue) 0.509 ± 0.282 (4e-03...1e+00)

To select multiple named slices, pass a tuple, list, or comma-separated string.

In [21]:
images.channels['red,blue']
Out[21]:
(setᵇ=4, xˢ=28, yˢ=28, channelsᶜ=red,blue) 0.503 ± 0.288 (2e-04...1e+00)

You can iterate along a dimension or unstack a tensor along a dimension.

In [22]:
for image in images.set:
    print(math.mean(image))

0.5032149
0.50452393
0.4983345
0.50876576
In [23]:
list(images.set)
Out[23]:
[(xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.503 ± 0.292 (4e-04...1e+00),
 (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.505 ± 0.287 (6e-04...1e+00),
 (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.498 ± 0.290 (1e-04...1e+00),
 (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.509 ± 0.284 (2e-04...1e+00)]

You can even convert named slices to a dict or use them as keyword arguments.

In [24]:
dict(images.set)
Out[24]:
{0: (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.503 ± 0.292 (4e-04...1e+00),
 1: (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.505 ± 0.287 (6e-04...1e+00),
 2: (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.498 ± 0.290 (1e-04...1e+00),
 3: (xˢ=28, yˢ=28, channelsᶜ=red,green,blue) 0.509 ± 0.284 (2e-04...1e+00)}
In [25]:
def color_transform(red, green, blue):
    print(f"Red: {math.mean(red)}, Green: {math.mean(green)}, Blue: {math.mean(blue)}")

color_transform(**dict(images.channels))

Red: (0.485, 0.500, 0.503, 0.500) along setᵇ, Green: (0.500, 0.507, 0.503, 0.511) along setᵇ, Blue: (0.525, 0.507, 0.489, 0.515) along setᵇ

Further Reading¶

Learn more about the dimension types and how to efficiently operate on tensors.

ΦML unifies data types as well and lets you set the floating point precision globally or by context.

While the dimensionality of neural networks must be specified during network creation, this is not the case for math functions. These automatically adapt to the number of spatial dimensions of the data that is passed in.

🌐 ΦML   •   📖 Documentation   •   🔗 API   •   ▶ Videos   •   Examples