mirror of
https://github.com/open-mmlab/mmcv.git
synced 2025-06-03 21:54:52 +08:00
2a3d2d4a35
* Reimplement cc_attention using pure pytorch * fix: avoid BC-Breaking * delete cc_attention related cpp and cuda files * delete cc_attention related lines in pybind.cpp * make out Tensor contiguous. * remove unneeded lines. * Update mmcv/ops/cc_attention.py Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com> * Update TestCrissCrossAttention * passing pre-commit * Update docstring of CrissCrossAttention * Update docstring of CrissCrossAttention * Update mmcv/ops/cc_attention.py Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com> * [docs]polish the docstring * [Docs] Polish the docstring Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com> Co-authored-by: Zaida Zhou <58739961+zhouzaida@users.noreply.github.com>
56 lines
1.5 KiB
Python
56 lines
1.5 KiB
Python
import numpy as np
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
|
|
class Loss(nn.Module):
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
def forward(self, input, target):
|
|
input = input.view(-1)
|
|
target = target.view(-1)
|
|
return torch.mean(input - target)
|
|
|
|
|
|
class TestCrissCrossAttention(object):
|
|
|
|
def test_cc_attention(self):
|
|
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
|
|
|
|
from mmcv.ops import CrissCrossAttention
|
|
loss_func = Loss()
|
|
|
|
input = np.fromfile(
|
|
'tests/data/for_ccattention/ccattention_input.bin',
|
|
dtype=np.float32)
|
|
output = np.fromfile(
|
|
'tests/data/for_ccattention/ccattention_output.bin',
|
|
dtype=np.float32)
|
|
input = input.reshape((1, 32, 45, 45))
|
|
output = output.reshape((1, 32, 45, 45))
|
|
label = torch.ones((1, 32, 45, 45))
|
|
|
|
input = torch.FloatTensor(input)
|
|
output = torch.FloatTensor(output)
|
|
|
|
input.requires_grad = True
|
|
|
|
shape = input.shape
|
|
channel = shape[1]
|
|
|
|
cca = CrissCrossAttention(channel)
|
|
cca.to(device)
|
|
input = input.to(device)
|
|
label = label.to(device)
|
|
cca.train()
|
|
test_output = cca(input)
|
|
test_loss = loss_func(test_output, label)
|
|
test_loss.backward()
|
|
test_output = test_output.detach().cpu().numpy()
|
|
output = output.numpy()
|
|
|
|
assert np.allclose(test_output, output)
|
|
assert test_output.shape == shape
|