def test_meshgrid_indexing_runtime_error(test_case):
with test_case.assertRaises(Exception) as context:
x1 = flow.ones((2), dtype=flow.float32, requires_grad=True)
x2 = flow.ones((2), dtype=flow.float32, requires_grad=True)
y = flow.meshgrid(x1, x2, indexing="ab")
test_case.assertTrue(
"meshgrid: indexing must be one of" in str(context.exception))
def test_meshgrid_tensors_scalar_runtime_error(test_case):
with test_case.assertRaises(Exception) as context:
x1 = flow.tensor([], dtype=flow.float32, requires_grad=True)
x2 = flow.ones((1, 2, 3), dtype=flow.float32, requires_grad=True)
y = flow.meshgrid(x1, x2)
test_case.assertTrue("Expected scalar or 1D tensor in the tensor list"
in str(context.exception))
def test_meshgrid_tensors_dtype_runtime_error(test_case):
with test_case.assertRaises(Exception) as context:
x1 = flow.ones((2), dtype=flow.float32, requires_grad=True)
x2 = flow.ones((2), dtype=flow.float16, requires_grad=True)
y = flow.meshgrid(x1, x2)
test_case.assertTrue(
"meshgrid expects all tensors to have the same dtype" in str(
context.exception))
def _test_meshgrid_forawd_scalar(test_case, device, indexing):
input1 = flow.tensor(np.array(1.0),
dtype=flow.float32,
device=flow.device(device))
input2 = flow.tensor(np.array(2.0),
dtype=flow.float32,
device=flow.device(device))
(np_x, np_y) = np.meshgrid(input1.numpy(),
input2.numpy(),
indexing=indexing)
(of_x, of_y) = flow.meshgrid(input1, input2, indexing=indexing)
test_case.assertTrue(np.allclose(of_x.numpy(), np_x, 0.0001, 0.0001))
def test_meshgrid_tensors_placement_runtime_error(test_case):
with test_case.assertRaises(Exception) as context:
x1 = flow.tensor(
[0.0, 1.0],
dtype=flow.float32,
placement=flow.placement("cpu", ranks=[0]),
sbp=[flow.sbp.broadcast],
)
x2 = flow.tensor(
[0.0, 1.0],
dtype=flow.float32,
placement=flow.placement("cpu", ranks=[0]),
sbp=[flow.sbp.broadcast],
).to_local()
y = flow.meshgrid(x1, x2)
test_case.assertTrue("meshgrid expects all tensors are global tensor"
in str(context.exception))
def __init__(
self,
dim,
window_size,
num_heads,
qkv_bias=True,
qk_scale=None,
attn_drop=0.0,
proj_drop=0.0,
):
super().__init__()
self.dim = dim
self.window_size = window_size # Wh, Ww
self.num_heads = num_heads
head_dim = dim // num_heads
self.scale = qk_scale or head_dim**-0.5
# define a parameter table of relative position bias
# Author zzk: we add trunc normal here！
self.relative_position_bias_table = nn.Parameter(
flow.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1),
num_heads)) # 2*Wh-1 * 2*Ww-1, nH
self.relative_position_bias_table.trunc_normal_(std=0.02)
# get pair-wise relative position index for each token inside the window
coords_h = flow.arange(self.window_size[0])
coords_w = flow.arange(self.window_size[1])
coords = flow.stack(flow.meshgrid(*[coords_h, coords_w])) # 2, Wh, Ww
coords_flatten = flow.flatten(coords, 1) # 2, Wh*Ww
relative_coords = (coords_flatten[:, :, None] -
coords_flatten[:, None, :]) # 2, Wh*Ww, Wh*Ww
relative_coords = relative_coords.permute(1, 2, 0) # Wh*Ww, Wh*Ww, 2
relative_coords[:, :,
0] += self.window_size[0] - 1 # shift to start from 0
relative_coords[:, :, 1] += self.window_size[1] - 1
relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww
self.register_buffer("relative_position_index",
relative_position_index)
self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
self.attn_drop = nn.Dropout(attn_drop)
self.proj = nn.Linear(dim, dim)
self.proj_drop = nn.Dropout(proj_drop)
self.softmax = nn.Softmax(dim=-1)
def _test_meshgrid_forawd_3tensor(test_case, device, indexing):
input1 = flow.tensor(np.array([1, 2, 3]),
dtype=flow.float32,
device=flow.device(device))
input2 = flow.tensor(np.array([4, 5, 6]),
dtype=flow.float32,
device=flow.device(device))
input3 = flow.tensor(np.array([7, 8, 9]),
dtype=flow.float32,
device=flow.device(device))
(np_x, np_y, np_z) = np.meshgrid(input1.numpy(),
input2.numpy(),
input3.numpy(),
indexing=indexing)
(of_x, of_y, of_z) = flow.meshgrid(input1,
input2,
input3,
indexing=indexing)
test_case.assertTrue(np.allclose(of_x.numpy(), np_x, 0.0001, 0.0001))
def test_meshgrid_tensors_size_runtime_error(test_case):
with test_case.assertRaises(Exception) as context:
y = flow.meshgrid([])
test_case.assertTrue("meshgrid expects a non-empty TensorList" in str(
context.exception))