def backward(self, grad_output):
grad_input1 = torch.zeros(self.input1.size())
if grad_output.is_cuda:
self.batchgrid = self.batchgrid.cuda()
grad_input1 = grad_input1.cuda()
#print('gradout:',grad_output.size())
grad_output_temp = grad_output.contiguous()
grad_output_view = grad_output_temp.view(-1, self.height*self.width, 2)
grad_output_view.contiguous()
grad_output_temp = torch.transpose(grad_output_view, 1, 2)
grad_output_temp.contiguous()
batchgrid_temp = self.batchgrid.view(-1, self.height*self.width, 3)
batchgrid_temp.contiguous()
grad_input1 = torch.baddbmm(grad_input1, grad_output_temp, batchgrid_temp)
return grad_input1
python类baddbmm()的实例源码
def SkipConnectLSTMCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
input = input.expand(4, *input.size()) if noise_in is None else input.unsqueeze(0) * noise_in
hx, cx = hidden
hx = torch.cat([hx, hidden_skip], dim=1)
hx = hx.expand(4, *hx.size()) if noise_hidden is None else hx.unsqueeze(0) * noise_hidden
gates = torch.baddbmm(b_ih.unsqueeze(1), input, w_ih) + torch.baddbmm(b_hh.unsqueeze(1), hx, w_hh)
ingate, forgetgate, cellgate, outgate = gates
ingate = F.sigmoid(ingate)
forgetgate = F.sigmoid(forgetgate)
cellgate = F.tanh(cellgate)
outgate = F.sigmoid(outgate)
cy = (forgetgate * cx) + (ingate * cellgate)
hy = outgate * F.tanh(cy)
return hy, cy
def SkipConnectGRUCell(input, hidden, hidden_skip, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
input = input.expand(3, *input.size()) if noise_in is None else input.unsqueeze(0) * noise_in
hx = torch.cat([hidden, hidden_skip], dim=1)
hx = hx.expand(3, *hx.size()) if noise_hidden is None else hx.unsqueeze(0) * noise_hidden
gi = torch.baddbmm(b_ih.unsqueeze(1), input, w_ih)
gh = torch.baddbmm(b_hh.unsqueeze(1), hx, w_hh)
i_r, i_i, i_n = gi
h_r, h_i, h_n = gh
resetgate = F.sigmoid(i_r + h_r)
inputgate = F.sigmoid(i_i + h_i)
newgate = F.tanh(i_n + resetgate * h_n)
hy = newgate + inputgate * (hidden - newgate)
return hy
def VarLSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
input = input.expand(4, *input.size()) if noise_in is None else input.unsqueeze(0) * noise_in
hx, cx = hidden
hx = hx.expand(4, *hx.size()) if noise_hidden is None else hx.unsqueeze(0) * noise_hidden
gates = torch.baddbmm(b_ih.unsqueeze(1), input, w_ih) + torch.baddbmm(b_hh.unsqueeze(1), hx, w_hh)
ingate, forgetgate, cellgate, outgate = gates
ingate = F.sigmoid(ingate)
forgetgate = F.sigmoid(forgetgate)
cellgate = F.tanh(cellgate)
outgate = F.sigmoid(outgate)
cy = (forgetgate * cx) + (ingate * cellgate)
hy = outgate * F.tanh(cy)
return hy, cy
def forward(self, add_batch, batch1, batch2):
self.save_for_backward(batch1, batch2)
output = self._get_output(add_batch)
return torch.baddbmm(output, self.alpha, add_batch, self.beta,
batch1, batch2)
def test_baddbmm(self):
num_batches = 10
M, N, O = 12, 8, 5
b1 = torch.randn(num_batches, M, N)
b2 = torch.randn(num_batches, N, O)
res = torch.bmm(b1, b2)
res2 = torch.Tensor().resize_as_(res).zero_()
res2.baddbmm_(b1,b2)
self.assertEqual(res2, res)
res2.baddbmm_(1,b1,b2)
self.assertEqual(res2, res*2)
res2.baddbmm_(1,.5,b1,b2)
self.assertEqual(res2, res*2.5)
res3 = torch.baddbmm(1,res2,0,b1,b2)
self.assertEqual(res3, res2)
res4 = torch.baddbmm(1,res2,.5,b1,b2)
self.assertEqual(res4, res*3)
res5 = torch.baddbmm(0,res2,1,b1,b2)
self.assertEqual(res5, res)
res6 = torch.baddbmm(.1,res2,.5,b1,b2)
self.assertEqual(res6, res2 * .1 + res * .5)
def forward(ctx, add_batch, batch1, batch2, alpha=1, beta=1, inplace=False):
ctx.alpha = alpha
ctx.beta = beta
ctx.save_for_backward(batch1, batch2)
output = _get_output(ctx, add_batch, inplace=inplace)
return torch.baddbmm(alpha, add_batch, beta,
batch1, batch2, out=output)
def test_functional_blas(self):
def compare(fn, *args):
unpacked_args = tuple(arg.data if isinstance(arg, Variable) else arg
for arg in args)
self.assertEqual(fn(*args).data, fn(*unpacked_args))
def test_blas_add(fn, x, y, z):
# Checks all signatures
compare(fn, x, y, z)
compare(fn, 0.5, x, y, z)
compare(fn, 0.5, x, 0.25, y, z)
def test_blas(fn, x, y):
compare(fn, x, y)
test_blas(torch.mm, Variable(torch.randn(2, 10)),
Variable(torch.randn(10, 4)))
test_blas_add(torch.addmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10, 4)))
test_blas(torch.bmm, Variable(torch.randn(4, 2, 10)),
Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.addbmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.baddbmm, Variable(torch.randn(4, 2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas(torch.mv, Variable(torch.randn(2, 10)),
Variable(torch.randn(10)))
test_blas_add(torch.addmv, Variable(torch.randn(2)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10)))
test_blas(torch.ger, Variable(torch.randn(5)),
Variable(torch.randn(6)))
test_blas_add(torch.addr, Variable(torch.randn(5, 6)),
Variable(torch.randn(5)), Variable(torch.randn(6)))
def test_baddbmm(self):
num_batches = 10
M, N, O = 12, 8, 5
b1 = torch.randn(num_batches, M, N)
b2 = torch.randn(num_batches, N, O)
res = torch.bmm(b1, b2)
res2 = torch.Tensor().resize_as_(res).zero_()
res2.baddbmm_(b1, b2)
self.assertEqual(res2, res)
res2.baddbmm_(1, b1, b2)
self.assertEqual(res2, res * 2)
res2.baddbmm_(1, .5, b1, b2)
self.assertEqual(res2, res * 2.5)
res3 = torch.baddbmm(1, res2, 0, b1, b2)
self.assertEqual(res3, res2)
res4 = torch.baddbmm(1, res2, .5, b1, b2)
self.assertEqual(res4, res * 3)
res5 = torch.baddbmm(0, res2, 1, b1, b2)
self.assertEqual(res5, res)
res6 = torch.baddbmm(.1, res2, .5, b1, b2)
self.assertEqual(res6, res2 * .1 + res * .5)
def forward(self, input, indices=None):
"""
Shape:
- target_batch :math:`(N, E, 1+N_r)`where `N = length, E = embedding size, N_r = noise ratio`
"""
if indices is None:
return super(IndexLinear, self).forward(input)
# the pytorch's [] operator BP can't correctly
input = input.unsqueeze(1)
target_batch = self.weight.index_select(0, indices.view(-1)).view(indices.size(0), indices.size(1), -1).transpose(1,2)
bias = self.bias.index_select(0, indices.view(-1)).view(indices.size(0), 1, indices.size(1))
out = torch.baddbmm(1, bias, 1, input, target_batch)
return out.squeeze()
def forward(ctx, add_batch, batch1, batch2, alpha=1, beta=1, inplace=False):
ctx.alpha = alpha
ctx.beta = beta
ctx.save_for_backward(batch1, batch2)
output = _get_output(ctx, add_batch, inplace=inplace)
return torch.baddbmm(alpha, add_batch, beta,
batch1, batch2, out=output)
def test_functional_blas(self):
def compare(fn, *args):
unpacked_args = tuple(arg.data if isinstance(arg, Variable) else arg
for arg in args)
self.assertEqual(fn(*args).data, fn(*unpacked_args))
def test_blas_add(fn, x, y, z):
# Checks all signatures
compare(fn, x, y, z)
compare(fn, 0.5, x, y, z)
compare(fn, 0.5, x, 0.25, y, z)
def test_blas(fn, x, y):
compare(fn, x, y)
test_blas(torch.mm, Variable(torch.randn(2, 10)),
Variable(torch.randn(10, 4)))
test_blas_add(torch.addmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10, 4)))
test_blas(torch.bmm, Variable(torch.randn(4, 2, 10)),
Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.addbmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.baddbmm, Variable(torch.randn(4, 2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas(torch.mv, Variable(torch.randn(2, 10)),
Variable(torch.randn(10)))
test_blas_add(torch.addmv, Variable(torch.randn(2)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10)))
test_blas(torch.ger, Variable(torch.randn(5)),
Variable(torch.randn(6)))
test_blas_add(torch.addr, Variable(torch.randn(5, 6)),
Variable(torch.randn(5)), Variable(torch.randn(6)))
def test_baddbmm(self):
num_batches = 10
M, N, O = 12, 8, 5
b1 = torch.randn(num_batches, M, N)
b2 = torch.randn(num_batches, N, O)
res = torch.bmm(b1, b2)
res2 = torch.Tensor().resize_as_(res).zero_()
res2.baddbmm_(b1, b2)
self.assertEqual(res2, res)
res2.baddbmm_(1, b1, b2)
self.assertEqual(res2, res * 2)
res2.baddbmm_(1, .5, b1, b2)
self.assertEqual(res2, res * 2.5)
res3 = torch.baddbmm(1, res2, 0, b1, b2)
self.assertEqual(res3, res2)
res4 = torch.baddbmm(1, res2, .5, b1, b2)
self.assertEqual(res4, res * 3)
res5 = torch.baddbmm(0, res2, 1, b1, b2)
self.assertEqual(res5, res)
res6 = torch.baddbmm(.1, res2, .5, b1, b2)
self.assertEqual(res6, res2 * .1 + res * .5)
def forward(ctx, add_batch, batch1, batch2, alpha=1, beta=1, inplace=False):
ctx.alpha = alpha
ctx.beta = beta
ctx.add_batch_size = add_batch.size()
ctx.save_for_backward(batch1, batch2)
output = _get_output(ctx, add_batch, inplace=inplace)
return torch.baddbmm(alpha, add_batch, beta,
batch1, batch2, out=output)
def test_baddbmm(self):
num_batches = 10
M, N, O = 12, 8, 5
b1 = torch.randn(num_batches, M, N)
b2 = torch.randn(num_batches, N, O)
res = torch.bmm(b1, b2)
res2 = torch.Tensor().resize_as_(res).zero_()
res2.baddbmm_(b1, b2)
self.assertEqual(res2, res)
res2.baddbmm_(1, b1, b2)
self.assertEqual(res2, res * 2)
res2.baddbmm_(1, .5, b1, b2)
self.assertEqual(res2, res * 2.5)
res3 = torch.baddbmm(1, res2, 0, b1, b2)
self.assertEqual(res3, res2)
res4 = torch.baddbmm(1, res2, .5, b1, b2)
self.assertEqual(res4, res * 3)
res5 = torch.baddbmm(0, res2, 1, b1, b2)
self.assertEqual(res5, res)
res6 = torch.baddbmm(.1, res2, .5, b1, b2)
self.assertEqual(res6, res2 * .1 + res * .5)
def _test_broadcast_fused_matmul(self, cast):
fns = ["baddbmm", "addbmm", "addmm", "addmv", "addr"]
for fn in fns:
batch_dim = random.randint(1, 8)
n_dim = random.randint(1, 8)
m_dim = random.randint(1, 8)
p_dim = random.randint(1, 8)
def dims_full_for_fn():
if fn == "baddbmm":
return ([batch_dim, n_dim, p_dim], [batch_dim, n_dim, m_dim], [batch_dim, m_dim, p_dim])
elif fn == "addbmm":
return ([n_dim, p_dim], [batch_dim, n_dim, m_dim], [batch_dim, m_dim, p_dim])
elif fn == "addmm":
return ([n_dim, p_dim], [n_dim, m_dim], [m_dim, p_dim])
elif fn == "addmv":
return ([n_dim], [n_dim, m_dim], [m_dim])
elif fn == "addr":
return ([n_dim, m_dim], [n_dim], [m_dim])
else:
raise AssertionError("unknown function")
(t0_dims_full, t1_dims, t2_dims) = dims_full_for_fn()
(t0_dims_small, _, _) = self._select_broadcastable_dims(t0_dims_full)
t0_small = cast(torch.randn(*t0_dims_small).float())
t1 = cast(torch.randn(*t1_dims).float())
t2 = cast(torch.randn(*t2_dims).float())
t0_full = cast(t0_small.expand(*t0_dims_full))
fntorch = getattr(torch, fn)
r0 = fntorch(t0_small, t1, t2)
r1 = fntorch(t0_full, t1, t2)
self.assertEqual(r0, r1)
def VarGRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, noise_in=None, noise_hidden=None):
input = input.expand(3, *input.size()) if noise_in is None else input.unsqueeze(0) * noise_in
hx = hidden.expand(3, *hidden.size()) if noise_hidden is None else hidden.unsqueeze(0) * noise_hidden
gi = torch.baddbmm(b_ih.unsqueeze(1), input, w_ih)
gh = torch.baddbmm(b_hh.unsqueeze(1), hx, w_hh)
i_r, i_i, i_n = gi
h_r, h_i, h_n = gh
resetgate = F.sigmoid(i_r + h_r)
inputgate = F.sigmoid(i_i + h_i)
newgate = F.tanh(i_n + resetgate * h_n)
hy = newgate + inputgate * (hidden - newgate)
return hy
def backward(self, grad_output):
grad_input1 = self.input1.new(self.input1.size()).zero_()
# if grad_output.is_cuda:
# self.batchgrid = self.batchgrid.cuda()
# grad_input1 = grad_input1.cuda()
grad_input1 = torch.baddbmm(grad_input1, torch.transpose(grad_output.view(-1, self.height*self.width, 2), 1,2), self.batchgrid.view(-1, self.height*self.width, 3))
return grad_input1
def test_baddbmm(self):
num_batches = 10
M, N, O = 12, 8, 5
b1 = torch.randn(num_batches, M, N)
b2 = torch.randn(num_batches, N, O)
res = torch.bmm(b1, b2)
res2 = torch.Tensor().resize_as_(res).zero_()
res2.baddbmm_(b1, b2)
self.assertEqual(res2, res)
res2.baddbmm_(1, b1, b2)
self.assertEqual(res2, res * 2)
res2.baddbmm_(1, .5, b1, b2)
self.assertEqual(res2, res * 2.5)
res3 = torch.baddbmm(1, res2, 0, b1, b2)
self.assertEqual(res3, res2)
res4 = torch.baddbmm(1, res2, .5, b1, b2)
self.assertEqual(res4, res * 3)
res5 = torch.baddbmm(0, res2, 1, b1, b2)
self.assertEqual(res5, res)
res6 = torch.baddbmm(.1, res2, .5, b1, b2)
self.assertEqual(res6, res2 * .1 + res * .5)
def _test_broadcast_fused_matmul(self, cast):
fns = ["baddbmm", "addbmm", "addmm", "addmv", "addr"]
for fn in fns:
batch_dim = random.randint(1, 8)
n_dim = random.randint(1, 8)
m_dim = random.randint(1, 8)
p_dim = random.randint(1, 8)
def dims_full_for_fn():
if fn == "baddbmm":
return ([batch_dim, n_dim, p_dim], [batch_dim, n_dim, m_dim], [batch_dim, m_dim, p_dim])
elif fn == "addbmm":
return ([n_dim, p_dim], [batch_dim, n_dim, m_dim], [batch_dim, m_dim, p_dim])
elif fn == "addmm":
return ([n_dim, p_dim], [n_dim, m_dim], [m_dim, p_dim])
elif fn == "addmv":
return ([n_dim], [n_dim, m_dim], [m_dim])
elif fn == "addr":
return ([n_dim, m_dim], [n_dim], [m_dim])
else:
raise AssertionError("unknown function")
(t0_dims_full, t1_dims, t2_dims) = dims_full_for_fn()
(t0_dims_small, _, _) = self._select_broadcastable_dims(t0_dims_full)
t0_small = cast(torch.randn(*t0_dims_small).float())
t1 = cast(torch.randn(*t1_dims).float())
t2 = cast(torch.randn(*t2_dims).float())
t0_full = cast(t0_small.expand(*t0_dims_full))
fntorch = getattr(torch, fn)
r0 = fntorch(t0_small, t1, t2)
r1 = fntorch(t0_full, t1, t2)
self.assertEqual(r0, r1)
def backward(self, grad_output):
grad_input1 = torch.zeros(self.input1.size())
if grad_output.is_cuda:
self.batchgrid = self.batchgrid.cuda()
grad_input1 = grad_input1.cuda()
#print('gradout:',grad_output.size())
grad_input1 = torch.baddbmm(grad_input1, torch.transpose(grad_output.view(-1, self.height*self.width, 2), 1,2), self.batchgrid.view(-1, self.height*self.width, 3))
#print(grad_input1)
return grad_input1*self.lr
def test_functional_blas(self):
def compare(fn, *args):
unpacked_args = tuple(arg.data if isinstance(arg, Variable) else arg
for arg in args)
unpacked_result = fn(*unpacked_args)
packed_result = fn(*args).data
# if non-Variable torch function returns a scalar, compare to scalar
if not torch.is_tensor(unpacked_result):
assert packed_result.dim() == 1
assert packed_result.nelement() == 1
packed_result = packed_result[0]
self.assertEqual(packed_result, unpacked_result)
def test_blas_add(fn, x, y, z):
# Checks all signatures
compare(fn, x, y, z)
compare(fn, 0.5, x, y, z)
compare(fn, 0.5, x, 0.25, y, z)
def test_blas(fn, x, y):
compare(fn, x, y)
test_blas(torch.mm, Variable(torch.randn(2, 10)),
Variable(torch.randn(10, 4)))
test_blas_add(torch.addmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10, 4)))
test_blas(torch.bmm, Variable(torch.randn(4, 2, 10)),
Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.addbmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.baddbmm, Variable(torch.randn(4, 2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas(torch.mv, Variable(torch.randn(2, 10)),
Variable(torch.randn(10)))
test_blas_add(torch.addmv, Variable(torch.randn(2)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10)))
test_blas(torch.ger, Variable(torch.randn(5)),
Variable(torch.randn(6)))
test_blas_add(torch.addr, Variable(torch.randn(5, 6)),
Variable(torch.randn(5)), Variable(torch.randn(6)))
test_blas(torch.matmul, Variable(torch.randn(6)), Variable(torch.randn(6)))
test_blas(torch.matmul, Variable(torch.randn(10, 4)), Variable(torch.randn(4)))
test_blas(torch.matmul, Variable(torch.randn(5)), Variable(torch.randn(5, 6)))
test_blas(torch.matmul, Variable(torch.randn(2, 10)), Variable(torch.randn(10, 4)))
test_blas(torch.matmul, Variable(torch.randn(5, 2, 10)), Variable(torch.randn(5, 10, 4)))
test_blas(torch.matmul, Variable(torch.randn(3, 5, 2, 10)), Variable(torch.randn(3, 5, 10, 4)))
test_blas(torch.matmul, Variable(torch.randn(3, 5, 2, 10)), Variable(torch.randn(10)))
test_blas(torch.matmul, Variable(torch.randn(10)), Variable(torch.randn(3, 5, 10, 4)))
def test_functional_blas(self):
def compare(fn, *args):
unpacked_args = tuple(arg.data if isinstance(arg, Variable) else arg
for arg in args)
unpacked_result = fn(*unpacked_args)
packed_result = fn(*args).data
# if non-Variable torch function returns a scalar, compare to scalar
if not torch.is_tensor(unpacked_result):
assert packed_result.dim() == 1
assert packed_result.nelement() == 1
packed_result = packed_result[0]
self.assertEqual(packed_result, unpacked_result)
def test_blas_add(fn, x, y, z):
# Checks all signatures
compare(fn, x, y, z)
compare(fn, 0.5, x, y, z)
compare(fn, 0.5, x, 0.25, y, z)
def test_blas(fn, x, y):
compare(fn, x, y)
test_blas(torch.mm, Variable(torch.randn(2, 10)),
Variable(torch.randn(10, 4)))
test_blas_add(torch.addmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10, 4)))
test_blas(torch.bmm, Variable(torch.randn(4, 2, 10)),
Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.addbmm, Variable(torch.randn(2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas_add(torch.baddbmm, Variable(torch.randn(4, 2, 4)),
Variable(torch.randn(4, 2, 10)), Variable(torch.randn(4, 10, 4)))
test_blas(torch.mv, Variable(torch.randn(2, 10)),
Variable(torch.randn(10)))
test_blas_add(torch.addmv, Variable(torch.randn(2)),
Variable(torch.randn(2, 10)), Variable(torch.randn(10)))
test_blas(torch.ger, Variable(torch.randn(5)),
Variable(torch.randn(6)))
test_blas_add(torch.addr, Variable(torch.randn(5, 6)),
Variable(torch.randn(5)), Variable(torch.randn(6)))
test_blas(torch.matmul, Variable(torch.randn(6)), Variable(torch.randn(6)))
test_blas(torch.matmul, Variable(torch.randn(10, 4)), Variable(torch.randn(4)))
test_blas(torch.matmul, Variable(torch.randn(5)), Variable(torch.randn(5, 6)))
test_blas(torch.matmul, Variable(torch.randn(2, 10)), Variable(torch.randn(10, 4)))
test_blas(torch.matmul, Variable(torch.randn(5, 2, 10)), Variable(torch.randn(5, 10, 4)))
test_blas(torch.matmul, Variable(torch.randn(3, 5, 2, 10)), Variable(torch.randn(3, 5, 10, 4)))
test_blas(torch.matmul, Variable(torch.randn(3, 5, 2, 10)), Variable(torch.randn(10)))
test_blas(torch.matmul, Variable(torch.randn(10)), Variable(torch.randn(3, 5, 10, 4)))
