This is an operator written in CUDA for PyTorch.
To compute the global contrast among each pixel. This work is inspired by Xinyu Zhang and referred to the paper from Mingming Cheng Global Contrast Based Salient Region Detection.
def forward(
input,
)
"""
Params:
------
input: float tensor, shape (B, C, W, H)
Returns:
------
output: float tensor, shape (B, 1, W, H)
"""
def backward(
grad,
input
):
"""
Params:
------
grad: float tensor, shape (B, 1, W, H)
input: float tensor, shape (B, C, W, H)
Returns:
------
d_input: float tensor, shape (B, C, W, H)
"""
./install.sh
./test.sh <loop_time>
>>>import torch
>>>from global_contrast import GlobalContrast
>>>x = torch.rand((20, 16, 336, 336)).cuda()
>>>model = GlobalContrast()
>>>y = model(x)
>>>y.size
tensor([20, 1, 336, 336], device='cuda:0')
forward(ms) | backward(ms) | |
---|---|---|
naive | - | - |
cuda | 53.693 | 179.787 |