481 lines
18 KiB
Python
481 lines
18 KiB
Python
import pickle
|
|
from collections import namedtuple
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from einops import EinopsError, rearrange, reduce
|
|
from einops.tests import FLOAT_REDUCTIONS as REDUCTIONS
|
|
from einops.tests import collect_test_backends, is_backend_tested
|
|
|
|
__author__ = "Alex Rogozhnikov"
|
|
|
|
testcase = namedtuple("testcase", ["pattern", "axes_lengths", "input_shape", "wrong_shapes"])
|
|
|
|
rearrangement_patterns = [
|
|
testcase(
|
|
"b c h w -> b (c h w)",
|
|
dict(c=20),
|
|
(10, 20, 30, 40),
|
|
[(), (10,), (10, 10, 10), (10, 21, 30, 40), [1, 20, 1, 1, 1]],
|
|
),
|
|
testcase(
|
|
"b c (h1 h2) (w1 w2) -> b (c h2 w2) h1 w1",
|
|
dict(h2=2, w2=2),
|
|
(10, 20, 30, 40),
|
|
[(), (1, 1, 1, 1), (1, 10, 3), ()],
|
|
),
|
|
testcase(
|
|
"b ... c -> c b ...",
|
|
dict(b=10),
|
|
(10, 20, 30),
|
|
[(), (10,), (5, 10)],
|
|
),
|
|
]
|
|
|
|
|
|
def test_rearrange_imperative():
|
|
for backend in collect_test_backends(symbolic=False, layers=True):
|
|
print("Test layer for ", backend.framework_name)
|
|
|
|
for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns:
|
|
x = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape)
|
|
result_numpy = rearrange(x, pattern, **axes_lengths)
|
|
layer = backend.layers().Rearrange(pattern, **axes_lengths)
|
|
for shape in wrong_shapes:
|
|
try:
|
|
layer(backend.from_numpy(np.zeros(shape, dtype="float32")))
|
|
except BaseException:
|
|
pass
|
|
else:
|
|
raise AssertionError("Failure expected")
|
|
|
|
# simple pickling / unpickling
|
|
layer2 = pickle.loads(pickle.dumps(layer))
|
|
result1 = backend.to_numpy(layer(backend.from_numpy(x)))
|
|
result2 = backend.to_numpy(layer2(backend.from_numpy(x)))
|
|
assert np.allclose(result_numpy, result1)
|
|
assert np.allclose(result1, result2)
|
|
|
|
just_sum = backend.layers().Reduce("...->", reduction="sum")
|
|
|
|
variable = backend.from_numpy(x)
|
|
result = just_sum(layer(variable))
|
|
|
|
result.backward()
|
|
assert np.allclose(backend.to_numpy(variable.grad), 1)
|
|
|
|
|
|
def test_rearrange_symbolic():
|
|
for backend in collect_test_backends(symbolic=True, layers=True):
|
|
print("Test layer for ", backend.framework_name)
|
|
|
|
for pattern, axes_lengths, input_shape, _wrong_shapes in rearrangement_patterns:
|
|
x = np.arange(np.prod(input_shape), dtype="float32").reshape(input_shape)
|
|
result_numpy = rearrange(x, pattern, **axes_lengths)
|
|
layer = backend.layers().Rearrange(pattern, **axes_lengths)
|
|
input_shape_of_nones = [None] * len(input_shape)
|
|
shapes = [input_shape, input_shape_of_nones]
|
|
|
|
for shape in shapes:
|
|
symbol = backend.create_symbol(shape)
|
|
eval_inputs = [(symbol, x)]
|
|
|
|
result_symbol1 = layer(symbol)
|
|
result1 = backend.eval_symbol(result_symbol1, eval_inputs)
|
|
assert np.allclose(result_numpy, result1)
|
|
|
|
layer2 = pickle.loads(pickle.dumps(layer))
|
|
result_symbol2 = layer2(symbol)
|
|
result2 = backend.eval_symbol(result_symbol2, eval_inputs)
|
|
assert np.allclose(result1, result2)
|
|
|
|
# now testing back-propagation
|
|
just_sum = backend.layers().Reduce("...->", reduction="sum")
|
|
|
|
result_sum1 = backend.eval_symbol(just_sum(result_symbol1), eval_inputs)
|
|
result_sum2 = np.sum(x)
|
|
|
|
assert np.allclose(result_sum1, result_sum2)
|
|
|
|
|
|
reduction_patterns = [
|
|
*rearrangement_patterns,
|
|
testcase("b c h w -> b ()", dict(b=10), (10, 20, 30, 40), [(10,), (10, 20, 30)]),
|
|
testcase("b c (h1 h2) (w1 w2) -> b c h1 w1", dict(h1=15, h2=2, w2=2), (10, 20, 30, 40), [(10, 20, 31, 40)]),
|
|
testcase("b ... c -> b", dict(b=10), (10, 20, 30, 40), [(10,), (11, 10)]),
|
|
]
|
|
|
|
|
|
def test_reduce_imperative():
|
|
for backend in collect_test_backends(symbolic=False, layers=True):
|
|
print("Test layer for ", backend.framework_name)
|
|
for reduction in REDUCTIONS:
|
|
for pattern, axes_lengths, input_shape, wrong_shapes in reduction_patterns:
|
|
print(backend, reduction, pattern, axes_lengths, input_shape, wrong_shapes)
|
|
x = np.arange(1, 1 + np.prod(input_shape), dtype="float32").reshape(input_shape)
|
|
x /= x.mean()
|
|
result_numpy = reduce(x, pattern, reduction, **axes_lengths)
|
|
layer = backend.layers().Reduce(pattern, reduction, **axes_lengths)
|
|
for shape in wrong_shapes:
|
|
try:
|
|
layer(backend.from_numpy(np.zeros(shape, dtype="float32")))
|
|
except BaseException:
|
|
pass
|
|
else:
|
|
raise AssertionError("Failure expected")
|
|
|
|
# simple pickling / unpickling
|
|
layer2 = pickle.loads(pickle.dumps(layer))
|
|
result1 = backend.to_numpy(layer(backend.from_numpy(x)))
|
|
result2 = backend.to_numpy(layer2(backend.from_numpy(x)))
|
|
assert np.allclose(result_numpy, result1)
|
|
assert np.allclose(result1, result2)
|
|
|
|
just_sum = backend.layers().Reduce("...->", reduction="sum")
|
|
|
|
variable = backend.from_numpy(x)
|
|
result = just_sum(layer(variable))
|
|
|
|
result.backward()
|
|
grad = backend.to_numpy(variable.grad)
|
|
if reduction == "sum":
|
|
assert np.allclose(grad, 1)
|
|
if reduction == "mean":
|
|
assert np.allclose(grad, grad.min())
|
|
if reduction in ["max", "min"]:
|
|
assert np.all(np.isin(grad, [0, 1]))
|
|
assert np.sum(grad) > 0.5
|
|
|
|
|
|
def test_reduce_symbolic():
|
|
for backend in collect_test_backends(symbolic=True, layers=True):
|
|
print("Test layer for ", backend.framework_name)
|
|
for reduction in REDUCTIONS:
|
|
for pattern, axes_lengths, input_shape, _wrong_shapes in reduction_patterns:
|
|
x = np.arange(1, 1 + np.prod(input_shape), dtype="float32").reshape(input_shape)
|
|
x /= x.mean()
|
|
result_numpy = reduce(x, pattern, reduction, **axes_lengths)
|
|
layer = backend.layers().Reduce(pattern, reduction, **axes_lengths)
|
|
input_shape_of_nones = [None] * len(input_shape)
|
|
shapes = [input_shape, input_shape_of_nones]
|
|
|
|
for shape in shapes:
|
|
symbol = backend.create_symbol(shape)
|
|
eval_inputs = [(symbol, x)]
|
|
|
|
result_symbol1 = layer(symbol)
|
|
result1 = backend.eval_symbol(result_symbol1, eval_inputs)
|
|
assert np.allclose(result_numpy, result1)
|
|
|
|
layer2 = pickle.loads(pickle.dumps(layer))
|
|
result_symbol2 = layer2(symbol)
|
|
result2 = backend.eval_symbol(result_symbol2, eval_inputs)
|
|
assert np.allclose(result1, result2)
|
|
|
|
|
|
def create_torch_model(use_reduce=False, add_scripted_layer=False):
|
|
if not is_backend_tested("torch"):
|
|
pytest.skip()
|
|
else:
|
|
import torch.jit
|
|
from torch.nn import Conv2d, Linear, MaxPool2d, ReLU, Sequential
|
|
|
|
from einops.layers.torch import EinMix, Rearrange, Reduce
|
|
|
|
return Sequential(
|
|
Conv2d(3, 6, kernel_size=(5, 5)),
|
|
Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2) if use_reduce else MaxPool2d(kernel_size=2),
|
|
Conv2d(6, 16, kernel_size=(5, 5)),
|
|
Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2),
|
|
torch.jit.script(Rearrange("b c h w -> b (c h w)"))
|
|
if add_scripted_layer
|
|
else Rearrange("b c h w -> b (c h w)"),
|
|
Linear(16 * 5 * 5, 120),
|
|
ReLU(),
|
|
Linear(120, 84),
|
|
ReLU(),
|
|
EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84),
|
|
EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84),
|
|
Linear(84, 10),
|
|
)
|
|
|
|
|
|
def test_torch_layer():
|
|
if not is_backend_tested("torch"):
|
|
pytest.skip()
|
|
else:
|
|
# checked that torch present
|
|
import torch
|
|
import torch.jit
|
|
|
|
model1 = create_torch_model(use_reduce=True)
|
|
model2 = create_torch_model(use_reduce=False)
|
|
input = torch.randn([10, 3, 32, 32])
|
|
# random models have different predictions
|
|
assert not torch.allclose(model1(input), model2(input))
|
|
model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict())))
|
|
assert torch.allclose(model1(input), model2(input))
|
|
|
|
# tracing (freezing)
|
|
model3 = torch.jit.trace(model2, example_inputs=input)
|
|
torch.testing.assert_close(model1(input), model3(input), atol=1e-3, rtol=1e-3)
|
|
torch.testing.assert_close(model1(input + 1), model3(input + 1), atol=1e-3, rtol=1e-3)
|
|
|
|
model4 = torch.jit.trace(model2, example_inputs=input)
|
|
torch.testing.assert_close(model1(input), model4(input), atol=1e-3, rtol=1e-3)
|
|
torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_torch_layers_scripting():
|
|
if not is_backend_tested("torch"):
|
|
pytest.skip()
|
|
else:
|
|
import torch
|
|
|
|
for script_layer in [False, True]:
|
|
model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer)
|
|
model2 = torch.jit.script(model1)
|
|
input = torch.randn([10, 3, 32, 32])
|
|
|
|
torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3)
|
|
|
|
|
|
def test_keras_layer():
|
|
rng = np.random.default_rng()
|
|
if not is_backend_tested("tensorflow"):
|
|
pytest.skip()
|
|
else:
|
|
import tensorflow as tf
|
|
|
|
if tf.__version__ < "2.16.":
|
|
# current implementation of layers follows new TF interface
|
|
pytest.skip()
|
|
from tensorflow.keras.layers import Conv2D as Conv2d
|
|
from tensorflow.keras.layers import Dense as Linear
|
|
from tensorflow.keras.layers import ReLU
|
|
from tensorflow.keras.models import Sequential
|
|
|
|
from einops.layers.keras import EinMix, Rearrange, Reduce, keras_custom_objects
|
|
|
|
def create_keras_model():
|
|
return Sequential(
|
|
[
|
|
Conv2d(6, kernel_size=5, input_shape=[32, 32, 3]),
|
|
Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2),
|
|
Conv2d(16, kernel_size=5),
|
|
Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2),
|
|
Rearrange("b c h w -> b (c h w)"),
|
|
Linear(120),
|
|
ReLU(),
|
|
Linear(84),
|
|
ReLU(),
|
|
EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84),
|
|
EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84),
|
|
Linear(10),
|
|
]
|
|
)
|
|
|
|
model1 = create_keras_model()
|
|
model2 = create_keras_model()
|
|
|
|
input = rng.normal(size=[10, 32, 32, 3]).astype("float32")
|
|
# two randomly init models should provide different outputs
|
|
assert not np.allclose(model1.predict_on_batch(input), model2.predict_on_batch(input))
|
|
|
|
# get some temp filename
|
|
tmp_model_filename = "/tmp/einops_tf_model.h5"
|
|
# save arch + weights
|
|
print("temp_path_keras1", tmp_model_filename)
|
|
tf.keras.models.save_model(model1, tmp_model_filename)
|
|
model3 = tf.keras.models.load_model(tmp_model_filename, custom_objects=keras_custom_objects)
|
|
|
|
np.testing.assert_allclose(model1.predict_on_batch(input), model3.predict_on_batch(input))
|
|
|
|
weight_filename = "/tmp/einops_tf_model.weights.h5"
|
|
# save arch as json
|
|
model4 = tf.keras.models.model_from_json(model1.to_json(), custom_objects=keras_custom_objects)
|
|
model1.save_weights(weight_filename)
|
|
model4.load_weights(weight_filename)
|
|
model2.load_weights(weight_filename)
|
|
# check that differently-inialized model receives same weights
|
|
np.testing.assert_allclose(model1.predict_on_batch(input), model2.predict_on_batch(input))
|
|
# ulimate test
|
|
# save-load architecture, and then load weights - should return same result
|
|
np.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input))
|
|
|
|
|
|
def test_flax_layers():
|
|
"""
|
|
One-off simple tests for Flax layers.
|
|
Unfortunately, Flax layers have a different interface from other layers.
|
|
"""
|
|
if not is_backend_tested("jax"):
|
|
pytest.skip()
|
|
else:
|
|
import flax
|
|
import jax
|
|
import jax.numpy as jnp
|
|
from flax import linen as nn
|
|
|
|
from einops.layers.flax import EinMix, Rearrange, Reduce
|
|
|
|
class NN(nn.Module):
|
|
@nn.compact
|
|
def __call__(self, x):
|
|
x = EinMix(
|
|
"b (h h2) (w w2) c -> b h w c_out", "h2 w2 c c_out", "c_out", sizes=dict(h2=2, w2=3, c=4, c_out=5)
|
|
)(x)
|
|
x = Rearrange("b h w c -> b (w h c)", sizes=dict(c=5))(x)
|
|
x = Reduce("b hwc -> b", "mean", dict(hwc=2 * 3 * 5))(x)
|
|
return x
|
|
|
|
model = NN()
|
|
fixed_input = jnp.ones([10, 2 * 2, 3 * 3, 4])
|
|
params = model.init(jax.random.PRNGKey(0), fixed_input)
|
|
|
|
def eval_at_point(params):
|
|
return jnp.linalg.norm(model.apply(params, fixed_input))
|
|
|
|
vandg = jax.value_and_grad(eval_at_point)
|
|
value0 = eval_at_point(params)
|
|
value1, grad1 = vandg(params)
|
|
assert jnp.allclose(value0, value1)
|
|
if jax.__version__ < "0.6.0":
|
|
tree_map = jax.tree_map
|
|
else:
|
|
tree_map = jax.tree.map
|
|
|
|
params2 = tree_map(lambda x1, x2: x1 - x2 * 0.001, params, grad1)
|
|
|
|
value2 = eval_at_point(params2)
|
|
assert value0 >= value2, (value0, value2)
|
|
|
|
# check serialization
|
|
fbytes = flax.serialization.to_bytes(params)
|
|
_loaded = flax.serialization.from_bytes(params, fbytes)
|
|
|
|
|
|
def test_einmix_decomposition():
|
|
"""
|
|
Testing that einmix correctly decomposes into smaller transformations.
|
|
"""
|
|
from einops.layers._einmix import _EinmixDebugger
|
|
|
|
mixin1 = _EinmixDebugger(
|
|
"a b c d e -> e d c b a",
|
|
weight_shape="d a b",
|
|
d=2, a=3, b=5,
|
|
) # fmt: off
|
|
assert mixin1.pre_reshape_pattern is None
|
|
assert mixin1.post_reshape_pattern is None
|
|
assert mixin1.einsum_pattern == "abcde,dab->edcba"
|
|
assert mixin1.saved_weight_shape == [2, 3, 5]
|
|
assert mixin1.saved_bias_shape is None
|
|
|
|
mixin2 = _EinmixDebugger(
|
|
"a b c d e -> e d c b a",
|
|
weight_shape="d a b",
|
|
bias_shape="a b c d e",
|
|
a=1, b=2, c=3, d=4, e=5,
|
|
) # fmt: off
|
|
assert mixin2.pre_reshape_pattern is None
|
|
assert mixin2.post_reshape_pattern is None
|
|
assert mixin2.einsum_pattern == "abcde,dab->edcba"
|
|
assert mixin2.saved_weight_shape == [4, 1, 2]
|
|
assert mixin2.saved_bias_shape == [5, 4, 3, 2, 1]
|
|
|
|
mixin3 = _EinmixDebugger(
|
|
"... -> ...",
|
|
weight_shape="",
|
|
bias_shape="",
|
|
) # fmt: off
|
|
assert mixin3.pre_reshape_pattern is None
|
|
assert mixin3.post_reshape_pattern is None
|
|
assert mixin3.einsum_pattern == "...,->..."
|
|
assert mixin3.saved_weight_shape == []
|
|
assert mixin3.saved_bias_shape == []
|
|
|
|
mixin4 = _EinmixDebugger(
|
|
"b a ... -> b c ...",
|
|
weight_shape="b a c",
|
|
a=1, b=2, c=3,
|
|
) # fmt: off
|
|
assert mixin4.pre_reshape_pattern is None
|
|
assert mixin4.post_reshape_pattern is None
|
|
assert mixin4.einsum_pattern == "ba...,bac->bc..."
|
|
assert mixin4.saved_weight_shape == [2, 1, 3]
|
|
assert mixin4.saved_bias_shape is None
|
|
|
|
mixin5 = _EinmixDebugger(
|
|
"(b a) ... -> b c (...)",
|
|
weight_shape="b a c",
|
|
a=1, b=2, c=3,
|
|
) # fmt: off
|
|
assert mixin5.pre_reshape_pattern == "(b a) ... -> b a ..."
|
|
assert mixin5.pre_reshape_lengths == dict(a=1, b=2)
|
|
assert mixin5.post_reshape_pattern == "b c ... -> b c (...)"
|
|
assert mixin5.einsum_pattern == "ba...,bac->bc..."
|
|
assert mixin5.saved_weight_shape == [2, 1, 3]
|
|
assert mixin5.saved_bias_shape is None
|
|
|
|
mixin6 = _EinmixDebugger(
|
|
"b ... (a c) -> b ... (a d)",
|
|
weight_shape="c d",
|
|
bias_shape="a d",
|
|
a=1, c=3, d=4,
|
|
) # fmt: off
|
|
assert mixin6.pre_reshape_pattern == "b ... (a c) -> b ... a c"
|
|
assert mixin6.pre_reshape_lengths == dict(a=1, c=3)
|
|
assert mixin6.post_reshape_pattern == "b ... a d -> b ... (a d)"
|
|
assert mixin6.einsum_pattern == "b...ac,cd->b...ad"
|
|
assert mixin6.saved_weight_shape == [3, 4]
|
|
assert mixin6.saved_bias_shape == [1, 1, 4] # (b) a d, ellipsis does not participate
|
|
|
|
mixin7 = _EinmixDebugger(
|
|
"a ... (b c) -> a (... d b)",
|
|
weight_shape="c d b",
|
|
bias_shape="d b",
|
|
b=2, c=3, d=4,
|
|
) # fmt: off
|
|
assert mixin7.pre_reshape_pattern == "a ... (b c) -> a ... b c"
|
|
assert mixin7.pre_reshape_lengths == dict(b=2, c=3)
|
|
assert mixin7.post_reshape_pattern == "a ... d b -> a (... d b)"
|
|
assert mixin7.einsum_pattern == "a...bc,cdb->a...db"
|
|
assert mixin7.saved_weight_shape == [3, 4, 2]
|
|
assert mixin7.saved_bias_shape == [1, 4, 2] # (a) d b, ellipsis does not participate
|
|
|
|
|
|
def test_einmix_restrictions():
|
|
"""
|
|
Testing different cases
|
|
"""
|
|
from einops.layers._einmix import _EinmixDebugger
|
|
|
|
with pytest.raises(EinopsError):
|
|
_EinmixDebugger(
|
|
"a b c d e -> e d c b a",
|
|
weight_shape="d a b",
|
|
d=2, a=3, # missing b
|
|
) # fmt: off
|
|
|
|
with pytest.raises(EinopsError):
|
|
_EinmixDebugger(
|
|
"a b c d e -> e d c b a",
|
|
weight_shape="w a b",
|
|
d=2, a=3, b=1 # missing d
|
|
) # fmt: off
|
|
|
|
with pytest.raises(EinopsError):
|
|
_EinmixDebugger(
|
|
"(...) a -> ... a",
|
|
weight_shape="a", a=1, # ellipsis on the left
|
|
) # fmt: off
|
|
|
|
with pytest.raises(EinopsError):
|
|
_EinmixDebugger(
|
|
"(...) a -> a ...",
|
|
weight_shape="a", a=1, # ellipsis on the right side after bias axis
|
|
bias_shape="a",
|
|
) # fmt: off
|