def backward(self, grad_output): if self._indices is not None: indices = self._indices else: indices, = self.saved_tensors grad_output = grad_output.contiguous() if not self.sparse: if indices.dim() == 2: indices = indices.view(-1) with torch.cuda.device_of(grad_output): if grad_output.is_cuda: _sorted = torch.cuda.LongTensor() _indices = torch.cuda.LongTensor() _count = torch.cuda.LongTensor() else: _count = torch.IntTensor() _sorted = _indices = None # TODO: sparse updates... grad_weight = grad_output.new(self._weight_size).zero_() self._backend.LookupTable_accGradParameters( self._backend.library_state, indices, grad_output, grad_weight, _count, _sorted, _indices, self.scale_grad_by_freq, self.padding_idx, 1) else: sp = self._make_sparse(indices, type(grad_output)) go = grad_output.view(-1, grad_output.size()[-1]) grad_weight = torch.smm(sp, go) return None, grad_weight
def sds_bmm_torch(s_t1, d_t2): """ bmm (Batch Matrix Matrix) for sparse x dense -> sparse. This function doesn't support gradient. And sparse tensors cannot accept gradient due to the limitation of torch implementation. with s_t1.shape = (b, x, s), d_t2.shape = (b, s, y), the output shape is (b, x, y) This is a work around utilizing torch.smm for sparse x dense -> sparse :param s_t1: sparse tensor 1 (in list, representing batches) :param d_t2: dense tensor 2 :return: bmm result in sparse (in list, representing batches) """ device = d_t2.device assert type(s_t1) == list batch_num = len(s_t1) assert batch_num == d_t2.shape[0], 'Batch size mismatch.' outp = [] for b in range(batch_num): # force cpu _s_t1 = s_t1[b].cpu() _d_t2 = d_t2[b].cpu() assert _s_t1.shape[1] == _d_t2.shape[0], 'Matrix shape mismatch.' _outp = torch.smm(_s_t1, _d_t2) # CUDA version of smm is not implemented outp.append(_outp) return outp
def test(self): texture_img = cv2.imread('models/default_texture2.jpg') texture_img = torch.from_numpy(texture_img).unsqueeze(0).float() texture_img = texture_img.reshape(1, -1).transpose(0, 1) start_time = time.time() action_tensor = random.choice(self.action_sparse_tensor_data)['mat'] result_flat = torch.smm(action_tensor, texture_img).to_dense() result_flat = result_flat.transpose(0, 1) result_flat = result_flat.reshape(1, 224, 224, 3) stop_time = time.time() print('time use: {}'.format(stop_time - start_time)) result_flat = result_flat.numpy()[0, :] cv2.imshow('result', result_flat.astype(np.uint8)) cv2.waitKey()
def test_shape(di, dj, dk): x = self._gen_sparse(2, 20, [di, dj])[0] t = self._gen_sparse(2, 20, [di, dk])[0] y = torch.randn(dj, dk) alpha = random.random() beta = random.random() res = torch.saddmm(alpha, t, beta, x, y) expected = torch.addmm(alpha, self.safeToDense(t), beta, self.safeToDense(x), y) self.assertEqual(self.safeToDense(res), expected) res = torch.saddmm(t, x, y) expected = torch.addmm(self.safeToDense(t), self.safeToDense(x), y) self.assertEqual(self.safeToDense(res), expected) res = torch.smm(x, y) expected = torch.mm(self.safeToDense(x), y) self.assertEqual(self.safeToDense(res), expected)
def test_shape(di, dj, dk): x = self._gen_sparse(2, 20, [di, dj])[0] t = self._gen_sparse(2, 20, [di, dk])[0] y = torch.randn(dj, dk) alpha = random.random() beta = random.random() expected = torch.addmm(alpha, t.to_dense(), beta, x.to_dense(), y) res = torch.saddmm(alpha, t, beta, x, y) self.assertEqual(res.to_dense(), expected) expected = torch.addmm(t.to_dense(), x.to_dense(), y) res = torch.saddmm(t, x, y) self.assertEqual(res.to_dense(), expected) expected = torch.mm(x.to_dense(), y) res = torch.smm(x, y) self.assertEqual(res.to_dense(), expected)