298 lines
11 KiB
Python
298 lines
11 KiB
Python
import numpy as np
|
|
import pytest
|
|
|
|
from einops import parse_shape, rearrange, reduce
|
|
from einops.tests import is_backend_tested
|
|
from einops.tests.test_ops import imp_op_backends
|
|
|
|
|
|
def test_rearrange_examples():
|
|
def test1(x):
|
|
# transpose
|
|
y = rearrange(x, "b c h w -> b h w c")
|
|
assert tuple(y.shape) == (10, 30, 40, 20)
|
|
return y
|
|
|
|
def test2(x):
|
|
# view / reshape
|
|
y = rearrange(x, "b c h w -> b (c h w)")
|
|
assert tuple(y.shape) == (10, 20 * 30 * 40)
|
|
return y
|
|
|
|
def test3(x):
|
|
# depth-to-space
|
|
y = rearrange(x, "b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2)
|
|
assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2)
|
|
return y
|
|
|
|
def test4(x):
|
|
# space-to-depth
|
|
y = rearrange(x, "b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2)
|
|
assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2)
|
|
return y
|
|
|
|
def test5(x):
|
|
# simple transposition
|
|
y = rearrange(x, "b1 sound b2 letter -> b1 b2 sound letter")
|
|
assert tuple(y.shape) == (10, 30, 20, 40)
|
|
return y
|
|
|
|
def test6(x):
|
|
# parsing parameters
|
|
t = rearrange(x, "b c h w -> (b h w) c")
|
|
t = t[:, ::2] # replacement for dot-product, just changes size of second axis
|
|
assert tuple(t.shape) == (10 * 30 * 40, 10)
|
|
|
|
y = rearrange(t, "(b h w) c2 -> b c2 h w", **parse_shape(x, "b _ h w"))
|
|
assert tuple(y.shape) == (10, 10, 30, 40)
|
|
return y
|
|
|
|
def test7(x):
|
|
# split of embedding into groups
|
|
y1, y2 = rearrange(x, "b (c g) h w -> g b c h w", g=2)
|
|
assert tuple(y1.shape) == (10, 10, 30, 40)
|
|
assert tuple(y2.shape) == (10, 10, 30, 40)
|
|
return y1 + y2 # only one tensor is expected in output
|
|
|
|
def test8(x):
|
|
# max-pooling
|
|
y = reduce(x, "b c (h h1) (w w1) -> b c h w", reduction="max", h1=2, w1=2)
|
|
assert tuple(y.shape) == (10, 20, 30 // 2, 40 // 2)
|
|
return y
|
|
|
|
def test9(x):
|
|
# squeeze - unsqueeze
|
|
y = reduce(x, "b c h w -> b c () ()", reduction="max")
|
|
assert tuple(y.shape) == (10, 20, 1, 1)
|
|
y = rearrange(y, "b c () () -> c b")
|
|
assert tuple(y.shape) == (20, 10)
|
|
return y
|
|
|
|
def test10(x):
|
|
# stack
|
|
tensors = list(x + 0) # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185
|
|
tensors = rearrange(tensors, "b c h w -> b h w c")
|
|
assert tuple(tensors.shape) == (10, 30, 40, 20)
|
|
return tensors
|
|
|
|
def test11(x):
|
|
# concatenate
|
|
tensors = list(x + 0) # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185
|
|
tensors = rearrange(tensors, "b c h w -> h (b w) c")
|
|
assert tuple(tensors.shape) == (30, 10 * 40, 20)
|
|
return tensors
|
|
|
|
def shufflenet(x, convolve, c1, c2):
|
|
# shufflenet reordering example
|
|
x = convolve(x)
|
|
x = rearrange(x, "b (c1 c2) h w-> b (c2 c1) h w", c1=c1, c2=c2)
|
|
x = convolve(x)
|
|
return x
|
|
|
|
def convolve_strided_1d(x, stride, usual_convolution):
|
|
x = rearrange(x, "b c t1 t2 -> b c (t1 t2)") # reduce dimensionality
|
|
x = rearrange(x, "b c (t stride) -> (stride b) c t", stride=stride)
|
|
x = usual_convolution(x)
|
|
x = rearrange(x, "(stride b) c t -> b c (t stride)", stride=stride)
|
|
return x
|
|
|
|
def convolve_strided_2d(x, h_stride, w_stride, usual_convolution):
|
|
x = rearrange(x, "b c (h hs) (w ws) -> (hs ws b) c h w", hs=h_stride, ws=w_stride)
|
|
x = usual_convolution(x)
|
|
x = rearrange(x, "(hs ws b) c h w -> b c (h hs) (w ws)", hs=h_stride, ws=w_stride)
|
|
return x
|
|
|
|
def unet_like_1d(x, usual_convolution):
|
|
# u-net like steps for increasing / reducing dimensionality
|
|
x = rearrange(x, "b c t1 t2 -> b c (t1 t2)") # reduce dimensionality
|
|
y = rearrange(x, "b c (t dt) -> b (dt c) t", dt=2)
|
|
y = usual_convolution(y)
|
|
x = x + rearrange(y, "b (dt c) t -> b c (t dt)", dt=2)
|
|
return x
|
|
|
|
# mock for convolution (works for all backends)
|
|
def convolve_mock(x):
|
|
return x
|
|
|
|
tests = [
|
|
test1,
|
|
test2,
|
|
test3,
|
|
test4,
|
|
test5,
|
|
test6,
|
|
test7,
|
|
test8,
|
|
test9,
|
|
test10,
|
|
test11,
|
|
lambda x: shufflenet(x, convolve=convolve_mock, c1=4, c2=5),
|
|
lambda x: convolve_strided_1d(x, stride=2, usual_convolution=convolve_mock),
|
|
lambda x: convolve_strided_2d(x, h_stride=2, w_stride=2, usual_convolution=convolve_mock),
|
|
lambda x: unet_like_1d(x, usual_convolution=convolve_mock),
|
|
]
|
|
|
|
for backend in imp_op_backends:
|
|
print("testing source_examples for ", backend.framework_name)
|
|
for test in tests:
|
|
x = np.arange(10 * 20 * 30 * 40).reshape([10, 20, 30, 40])
|
|
result1 = test(x)
|
|
result2 = backend.to_numpy(test(backend.from_numpy(x)))
|
|
assert np.array_equal(result1, result2)
|
|
|
|
# now with strides
|
|
x = np.arange(10 * 2 * 20 * 3 * 30 * 1 * 40).reshape([10 * 2, 20 * 3, 30 * 1, 40 * 1])
|
|
# known torch bug - torch doesn't support negative steps
|
|
last_step = -1 if (backend.framework_name != "torch" and backend.framework_name != "oneflow") else 1
|
|
indexing_expression = np.index_exp[::2, ::3, ::1, ::last_step]
|
|
result1 = test(x[indexing_expression])
|
|
result2 = backend.to_numpy(test(backend.from_numpy(x)[indexing_expression]))
|
|
assert np.array_equal(result1, result2)
|
|
|
|
|
|
def tensor_train_example_numpy():
|
|
# kept here just for a collection, only tested for numpy
|
|
# https://arxiv.org/pdf/1509.06569.pdf, (5)
|
|
x = np.ones([3, 4, 5, 6])
|
|
rank = 4
|
|
if np.__version__ < "1.15.0":
|
|
# numpy.einsum fails here, skip test
|
|
return
|
|
# creating appropriate Gs
|
|
Gs = [np.ones([d, d, rank, rank]) for d in x.shape]
|
|
Gs[0] = Gs[0][:, :, :1, :]
|
|
Gs[-1] = Gs[-1][:, :, :, :1]
|
|
|
|
# einsum way
|
|
y = x.reshape((1, *x.shape))
|
|
for G in Gs:
|
|
# taking partial results left-to-right
|
|
# y = numpy.einsum('i j alpha beta, alpha i ... -> beta ... j', G, y)
|
|
y = np.einsum("i j a b, a i ... -> b ... j", G, y)
|
|
y1 = y.reshape(-1)
|
|
|
|
# alternative way
|
|
y = x.reshape(-1)
|
|
for G in Gs:
|
|
i, j, alpha, beta = G.shape
|
|
y = rearrange(y, "(i rest alpha) -> rest (alpha i)", alpha=alpha, i=i)
|
|
y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)")
|
|
y = rearrange(y, "rest (beta j) -> (beta rest j)", beta=beta, j=j)
|
|
y2 = y
|
|
assert np.allclose(y1, y2)
|
|
|
|
# yet another way
|
|
y = x
|
|
for G in Gs:
|
|
i, j, alpha, beta = G.shape
|
|
y = rearrange(y, "i ... (j alpha) -> ... j (alpha i)", alpha=alpha, i=i)
|
|
y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)")
|
|
y3 = y.reshape(-1)
|
|
assert np.allclose(y1, y3)
|
|
|
|
|
|
def test_pytorch_yolo_fragment():
|
|
if not is_backend_tested("torch"):
|
|
pytest.skip()
|
|
|
|
import torch
|
|
|
|
def old_way(tensor, num_classes, num_anchors, anchors, stride_h, stride_w):
|
|
# https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44
|
|
bs = tensor.size(0)
|
|
in_h = tensor.size(2)
|
|
in_w = tensor.size(3)
|
|
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]
|
|
|
|
prediction = tensor.view(bs, num_anchors, 5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous()
|
|
# Get outputs
|
|
x = torch.sigmoid(prediction[..., 0]) # Center x
|
|
y = torch.sigmoid(prediction[..., 1]) # Center y
|
|
w = prediction[..., 2] # Width
|
|
h = prediction[..., 3] # Height
|
|
conf = torch.sigmoid(prediction[..., 4]) # Conf
|
|
pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred.
|
|
|
|
# https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L70-L92
|
|
FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
|
|
LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
|
|
# Calculate offsets for each grid
|
|
grid_x = (
|
|
torch.linspace(0, in_w - 1, in_w)
|
|
.repeat(in_w, 1)
|
|
.repeat(bs * num_anchors, 1, 1)
|
|
.view(x.shape)
|
|
.type(FloatTensor)
|
|
)
|
|
grid_y = (
|
|
torch.linspace(0, in_h - 1, in_h)
|
|
.repeat(in_h, 1)
|
|
.t()
|
|
.repeat(bs * num_anchors, 1, 1)
|
|
.view(y.shape)
|
|
.type(FloatTensor)
|
|
)
|
|
# Calculate anchor w, h
|
|
anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0]))
|
|
anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1]))
|
|
anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape)
|
|
anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape)
|
|
# Add offset and scale with anchors
|
|
pred_boxes = FloatTensor(prediction[..., :4].shape)
|
|
pred_boxes[..., 0] = x.data + grid_x
|
|
pred_boxes[..., 1] = y.data + grid_y
|
|
pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
|
|
pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
|
|
# Results
|
|
_scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor)
|
|
output = torch.cat(
|
|
(pred_boxes.view(bs, -1, 4) * _scale, conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1
|
|
)
|
|
return output
|
|
|
|
def new_way(tensor, num_classes, num_anchors, anchors, stride_h, stride_w):
|
|
raw_predictions = rearrange(tensor, " b (anchor prediction) h w -> prediction b anchor h w", anchor=num_anchors)
|
|
|
|
anchors = torch.FloatTensor(anchors).to(tensor.device)
|
|
anchor_sizes = rearrange(anchors, "anchor dim -> dim () anchor () ()")
|
|
|
|
_, _, _, in_h, in_w = raw_predictions.shape
|
|
grid_h = rearrange(torch.arange(in_h).float(), "h -> () () h ()").to(tensor.device)
|
|
grid_w = rearrange(torch.arange(in_w).float(), "w -> () () () w").to(tensor.device)
|
|
|
|
predicted_bboxes = torch.zeros_like(raw_predictions)
|
|
predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h # center y
|
|
predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w # center x
|
|
predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes # bbox width and height
|
|
predicted_bboxes[4] = raw_predictions[4].sigmoid() # confidence
|
|
predicted_bboxes[5:] = raw_predictions[5:].sigmoid() # class predictions
|
|
# only to match results of original code, not needed
|
|
return rearrange(predicted_bboxes, "prediction b anchor h w -> b anchor h w prediction")
|
|
|
|
stride_h = 4
|
|
stride_w = 4
|
|
batch_size = 5
|
|
num_classes = 12
|
|
anchors = [[50, 100], [100, 50], [75, 75]]
|
|
num_anchors = len(anchors)
|
|
|
|
x = torch.randn([batch_size, num_anchors * (5 + num_classes), 1, 1])
|
|
result1 = old_way(
|
|
tensor=x,
|
|
num_anchors=num_anchors,
|
|
num_classes=num_classes,
|
|
stride_h=stride_h,
|
|
stride_w=stride_w,
|
|
anchors=anchors,
|
|
)
|
|
result2 = new_way(
|
|
tensor=x,
|
|
num_anchors=num_anchors,
|
|
num_classes=num_classes,
|
|
stride_h=stride_h,
|
|
stride_w=stride_w,
|
|
anchors=anchors,
|
|
)
|
|
result1 = result1.reshape(result2.shape)
|
|
assert torch.allclose(result1, result2)
|