add LR-Net V2
This commit is contained in:
v-zeliu1
parent
3dc2a55301
commit
be50b6cc51
108
README.md
108
README.md
|
|
@ -1,92 +1,24 @@
|
|||
# Swin Transformer
|
||||
# Local Relation Networks V2 (LR-Net V2)
|
||||
|
||||
[](https://paperswithcode.com/sota/object-detection-on-coco?p=swin-transformer-hierarchical-vision)
|
||||
[](https://paperswithcode.com/sota/instance-segmentation-on-coco?p=swin-transformer-hierarchical-vision)
|
||||
[](https://paperswithcode.com/sota/object-detection-on-coco-minival?p=swin-transformer-hierarchical-vision)
|
||||
[](https://paperswithcode.com/sota/instance-segmentation-on-coco-minival?p=swin-transformer-hierarchical-vision)
|
||||
[](https://paperswithcode.com/sota/semantic-segmentation-on-ade20k?p=swin-transformer-hierarchical-vision)
|
||||
[](https://paperswithcode.com/sota/semantic-segmentation-on-ade20k-val?p=swin-transformer-hierarchical-vision)
|
||||
This branch is an improved implementation of ["Local Relation Networks for Image Recognition (LR-Net)"](https://arxiv.org/pdf/1904.11491.pdf). The original LR-Net utilizes sliding window based self-attention layer to replace the `3x3` convolution layers in a ResNet architecture. This improved implementation applies this layer into a stronger overall architecture based on Tranformers, dubbed as LR-Net V2. We provide cuda kernels for the local relation layers. Training scripts and pre-trained models will be provided in the future.
|
||||
|
||||
By [Ze Liu](https://github.com/zeliu98/)\*, [Yutong Lin](https://github.com/impiga)\*, [Yue Cao](http://yue-cao.me)\*, [Han Hu](https://ancientmooner.github.io/)\*, [Yixuan Wei](https://github.com/weiyx16), [Zheng Zhang](https://stupidzz.github.io/), [Stephen Lin](https://scholar.google.com/citations?user=c3PYmxUAAAAJ&hl=en) and [Baining Guo](https://www.microsoft.com/en-us/research/people/bainguo/).
|
||||
## Install
|
||||
```bash
|
||||
cd ops/local_relation
|
||||
python setup.py build_ext --inplace
|
||||
```
|
||||
|
||||
This repo is the official implementation of ["Swin Transformer: Hierarchical Vision Transformer using Shifted Windows"](https://arxiv.org/pdf/2103.14030.pdf). It currently includes code and models for the following tasks:
|
||||
|
||||
> **Image Classification**: Included in this repo. See [get_started.md](get_started.md) for a quick start.
|
||||
|
||||
> **Object Detection and Instance Segmentation**: See [Swin Transformer for Object Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection).
|
||||
|
||||
> **Semantic Segmentation**: See [Swin Transformer for Semantic Segmentation](https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation).
|
||||
|
||||
## Updates
|
||||
|
||||
***04/12/2021***
|
||||
|
||||
Initial commits:
|
||||
|
||||
1. Pretrained models on ImageNet-1K ([Swin-T-IN1K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth), [Swin-S-IN1K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth), [Swin-B-IN1K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth)) and ImageNet-22K ([Swin-B-IN22K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth), [Swin-L-IN22K](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth)) are provided.
|
||||
2. The supported code and models for ImageNet-1K image classification, COCO object detection and ADE20K semantic segmentation are provided.
|
||||
3. The cuda kernel implementation for the [local relation layer](https://arxiv.org/pdf/1904.11491.pdf) is provided in branch [LR-Net](https://github.com/microsoft/Swin-Transformer/tree/LR-Net).
|
||||
|
||||
## Introduction
|
||||
|
||||
**Swin Transformer** is initially described in [arxiv](https://arxiv.org/abs/2103.14030), which capably serves as a
|
||||
general-purpose backbone for computer vision. It is basically a hierarchical Transformer whose representation is
|
||||
computed with shifted windows. The shifted windowing scheme brings greater efficiency by limiting self-attention
|
||||
computation to non-overlapping local windows while also allowing for cross-window connection.
|
||||
|
||||
Swin Transformer achieves strong performance on COCO object detection (`58.7 box AP` and `51.1 mask AP` on test-dev) and
|
||||
ADE20K semantic segmentatiion (`53.5 mIoU` on val), surpassing previous models by a large margin.
|
||||
|
||||
|
||||
|
||||
## Main Results on ImageNet with Pretrained Models
|
||||
|
||||
**ImageNet-1K and ImageNet-22K Pretrained Models**
|
||||
|
||||
| name | pretrain | resolution |acc@1 | acc@5 | #params | FLOPs | FPS| 22K model | 1K model |
|
||||
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |:---: |:---: |
|
||||
| Swin-T | ImageNet-1K | 224x224 | 81.2 | 95.5 | 28M | 4.5G | 755 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/156nWJy4Q28rDlrX-rRbI3w) |
|
||||
| Swin-S | ImageNet-1K | 224x224 | 83.2 | 96.2 | 50M | 8.7G | 437 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_small_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/1KFjpj3Efey3LmtE1QqPeQg) |
|
||||
| Swin-B | ImageNet-1K | 224x224 | 83.5 | 96.5 | 88M | 15.4G | 278 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224.pth)/[baidu](https://pan.baidu.com/s/16bqCTEc70nC_isSsgBSaqQ) |
|
||||
| Swin-B | ImageNet-1K | 384x384 | 84.5 | 97.0 | 88M | 47.1G | 85 | - | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384.pth)/[baidu](https://pan.baidu.com/s/1xT1cu740-ejW7htUdVLnmw) |
|
||||
| Swin-B | ImageNet-22K | 224x224 | 85.2 | 97.5 | 88M | 15.4G | 278 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22k.pth)/[baidu](https://pan.baidu.com/s/1y1Ec3UlrKSI8IMtEs-oBXA) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window7_224_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1n_wNkcbRxVXit8r_KrfAVg) |
|
||||
| Swin-B | ImageNet-22K | 384x384 | 86.4 | 98.0 | 88M | 47.1G | 85 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22k.pth)/[baidu](https://pan.baidu.com/s/1vwJxnJcVqcLZAw9HaqiR6g) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_base_patch4_window12_384_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1caKTSdoLJYoi4WBcnmWuWg) |
|
||||
| Swin-L | ImageNet-22K | 224x224 | 86.3 | 97.9 | 197M | 34.5G | 141 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22k.pth)/[baidu](https://pan.baidu.com/s/1pws3rOTFuOebBYP3h6Kx8w) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window7_224_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1NkQApMWUhxBGjk1ne6VqBQ) |
|
||||
| Swin-L | ImageNet-22K | 384x384 | 87.3 | 98.2 | 197M | 103.9G | 42 | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22k.pth)/[baidu](https://pan.baidu.com/s/1sl7o_bJA143OD7UqSLAMoA) | [github](https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_large_patch4_window12_384_22kto1k.pth)/[baidu](https://pan.baidu.com/s/1X0FLHQyPOC6Kmv2CmgxJvA) |
|
||||
|
||||
Note: access code for `baidu` is `swin`.
|
||||
|
||||
## Main Results on Downstream Tasks
|
||||
|
||||
**COCO Object Detection (2017 val)**
|
||||
|
||||
| Backbone | Method | pretrain | Lr Schd | box mAP | mask mAP | #params | FLOPs |
|
||||
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
||||
| Swin-T | Mask R-CNN | ImageNet-1K | 3x | 46.0 | 41.6 | 48M | 267G |
|
||||
| Swin-S | Mask R-CNN | ImageNet-1K | 3x | 48.5 | 43.3 | 69M | 359G |
|
||||
| Swin-T | Cascade Mask R-CNN | ImageNet-1K | 3x | 50.4 | 43.7 | 86M | 745G |
|
||||
| Swin-S | Cascade Mask R-CNN | ImageNet-1K | 3x | 51.9 | 45.0 | 107M | 838G |
|
||||
| Swin-B | Cascade Mask R-CNN | ImageNet-1K | 3x | 51.9 | 45.0 | 145M | 982G |
|
||||
| Swin-T | RepPoints V2 | ImageNet-1K | 3x | 50.0 | - | 45M | 283G |
|
||||
| Swin-T | Mask RepPoints V2 | ImageNet-1K | 3x | 50.3 | 43.6 | 47M | 292G |
|
||||
| Swin-B | HTC++ | ImageNet-22K | 6x | 56.4 | 49.1 | 160M | 1043G |
|
||||
| Swin-L | HTC++ | ImageNet-22K | 3x | 57.1 | 49.5 | 284M | 1470G |
|
||||
| Swin-L | HTC++<sup>*</sup> | ImageNet-22K | 3x | 58.0 | 50.4 | 284M | - |
|
||||
|
||||
Note: <sup>*</sup> indicates multi-scale testing.
|
||||
|
||||
**ADE20K Semantic Segmentation (val)**
|
||||
|
||||
| Backbone | Method | pretrain | Crop Size | Lr Schd | mIoU | mIoU (ms+flip) | #params | FLOPs |
|
||||
| :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: |
|
||||
| Swin-T | UPerNet | ImageNet-1K | 512x512 | 160K | 44.51 | 45.81 | 60M | 945G |
|
||||
| Swin-S | UperNet | ImageNet-1K | 512x512 | 160K | 47.64 | 49.47 | 81M | 1038G |
|
||||
| Swin-B | UperNet | ImageNet-1K | 512x512 | 160K | 48.13 | 49.72 | 121M | 1188G |
|
||||
| Swin-B | UPerNet | ImageNet-22K | 640x640 | 160K | 50.04 | 51.66 | 121M | 1841G |
|
||||
| Swin-L | UperNet | ImageNet-22K | 640x640 | 160K | 52.05 | 53.53 | 234M | 3230G |
|
||||
|
||||
## Citing Swin Transformer
|
||||
## Citing Local Relation Networks
|
||||
|
||||
```
|
||||
@inproceedings{hu2019local,
|
||||
title={Local relation networks for image recognition},
|
||||
author={Hu, Han and Zhang, Zheng and Xie, Zhenda and Lin, Stephen},
|
||||
booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision},
|
||||
pages={3464--3473},
|
||||
year={2019}
|
||||
}
|
||||
```
|
||||
```
|
||||
@article{liu2021Swin,
|
||||
title={Swin Transformer: Hierarchical Vision Transformer using Shifted Windows},
|
||||
|
|
@ -96,12 +28,6 @@ Note: <sup>*</sup> indicates multi-scale testing.
|
|||
}
|
||||
```
|
||||
|
||||
## Getting Started
|
||||
|
||||
- For **Image Classification**, please see [get_started.md](get_started.md) for detailed instructions.
|
||||
- For **Object Detection and Instance Segmentation**, please see [Swin Transformer for Object Detection](https://github.com/SwinTransformer/Swin-Transformer-Object-Detection).
|
||||
- For **Semantic Segmentation**, please see [Swin Transformer for Semantic Segmentation](https://github.com/SwinTransformer/Swin-Transformer-Semantic-Segmentation).
|
||||
|
||||
## Contributing
|
||||
|
||||
This project welcomes contributions and suggestions. Most contributions require you to agree to a
|
||||
|
|
|
|||
4
ops/local_relation/.gitignore
vendored
Normal file
4
ops/local_relation/.gitignore
vendored
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
*.so
|
||||
_ext*
|
||||
__pycache__
|
||||
build
|
||||
1
ops/local_relation/__init__.py
Normal file
1
ops/local_relation/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .local_relation_func import local_relation
|
||||
102
ops/local_relation/local_relation_func.py
Normal file
102
ops/local_relation/local_relation_func.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
# --------------------------------------------------------
|
||||
# Swin Transformer
|
||||
# Copyright (c) 2019 Microsoft
|
||||
# Licensed under The MIT License [see LICENSE for details]
|
||||
# Written by Han Hu, Jiarui Xu
|
||||
# Modified by Ze Liu
|
||||
# --------------------------------------------------------
|
||||
|
||||
import torch
|
||||
from torch.autograd import Function
|
||||
from torch.nn.modules.utils import _pair
|
||||
|
||||
from . import local_relation_cuda
|
||||
|
||||
|
||||
class LocalRelationFunction(Function):
|
||||
|
||||
@staticmethod
|
||||
def forward(ctx,
|
||||
query,
|
||||
key,
|
||||
value,
|
||||
pos_weight,
|
||||
kernel_size,
|
||||
groups,
|
||||
stride=1,
|
||||
dilation=1,
|
||||
scale=1.,
|
||||
no_define_value=-100.,
|
||||
norm_method=0,
|
||||
sim_method=0,
|
||||
batch_step=32):
|
||||
for input in [query, key, value]:
|
||||
if input is not None and input.dim() != 4:
|
||||
raise ValueError(
|
||||
"Expected 4D tensor as input, got {}D tensor instead.".format(
|
||||
input.dim()))
|
||||
ctx.kernel_size = _pair(kernel_size)
|
||||
ctx.groups = groups
|
||||
ctx.stride = stride
|
||||
ctx.dilation = dilation
|
||||
ctx.scale = scale
|
||||
ctx.no_define_value = no_define_value
|
||||
ctx.norm_method = norm_method
|
||||
ctx.sim_method = sim_method
|
||||
ctx.batch_step = batch_step
|
||||
|
||||
ctx.save_for_backward(query, key, value, pos_weight)
|
||||
|
||||
output = query.new_empty(
|
||||
LocalRelationFunction._output_size(query, value))
|
||||
|
||||
scale_tensor = query.new_tensor([ctx.scale])
|
||||
no_define_value_tensor = query.new_tensor([ctx.no_define_value])
|
||||
|
||||
if not input.is_cuda:
|
||||
raise NotImplementedError
|
||||
else:
|
||||
batch_step = min(ctx.batch_step, query.shape[0])
|
||||
local_relation_cuda.local_relation_forward_cuda(
|
||||
query, key, value, pos_weight, scale_tensor, no_define_value_tensor,
|
||||
output, ctx.kernel_size[1], ctx.kernel_size[0], ctx.groups,
|
||||
ctx.dilation, ctx.stride, batch_step, ctx.norm_method, ctx.sim_method)
|
||||
return output
|
||||
|
||||
@staticmethod
|
||||
def backward(ctx, grad_output):
|
||||
query, key, value, pos_weight = ctx.saved_tensors
|
||||
|
||||
grad_query = grad_key = grad_value = grad_pos_weight = None
|
||||
|
||||
scale_tensor = query.new_tensor(ctx.scale)
|
||||
no_define_value_tensor = query.new_tensor(ctx.no_define_value)
|
||||
|
||||
if not grad_output.is_cuda:
|
||||
raise NotImplementedError
|
||||
else:
|
||||
batch_step = min(ctx.batch_step, query.shape[0])
|
||||
|
||||
if ctx.needs_input_grad[0] or ctx.needs_input_grad[1] or ctx.needs_input_grad[2] or ctx.needs_input_grad[3]:
|
||||
grad_query = torch.zeros_like(query)
|
||||
grad_key = torch.zeros_like(key)
|
||||
grad_value = torch.zeros_like(value)
|
||||
grad_pos_weight = torch.zeros_like(pos_weight)
|
||||
local_relation_cuda.local_relation_backward_cuda(
|
||||
query, key, value, pos_weight,
|
||||
scale_tensor, no_define_value_tensor, grad_output,
|
||||
grad_query, grad_key, grad_value, grad_pos_weight,
|
||||
ctx.kernel_size[1], ctx.kernel_size[0],
|
||||
ctx.groups, ctx.dilation, ctx.stride, batch_step,
|
||||
ctx.norm_method, ctx.sim_method)
|
||||
|
||||
return (grad_query, grad_key, grad_value, grad_pos_weight, None, None, None,
|
||||
None, None, None, None, None, None)
|
||||
|
||||
@staticmethod
|
||||
def _output_size(query, value):
|
||||
output_size = (query.size(0), value.size(1), query.size(2), query.size(3))
|
||||
return output_size
|
||||
|
||||
|
||||
local_relation = LocalRelationFunction.apply
|
||||
12
ops/local_relation/setup.py
Normal file
12
ops/local_relation/setup.py
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
|
||||
setup(
|
||||
name='local_relation',
|
||||
ext_modules=[
|
||||
CUDAExtension('local_relation_cuda', [
|
||||
'src/local_relation_cuda.cpp',
|
||||
'src/local_relation_cuda_kernel.cu',
|
||||
]),
|
||||
],
|
||||
cmdclass={'build_ext': BuildExtension})
|
||||
331
ops/local_relation/src/local_relation_cuda.cpp
Normal file
331
ops/local_relation/src/local_relation_cuda.cpp
Normal file
|
|
@ -0,0 +1,331 @@
|
|||
/*!
|
||||
* Copyright (c) 2019 Microsoft
|
||||
* Licensed under The MIT License [see LICENSE for details]
|
||||
* \file local_relation_cuda.cpp
|
||||
* \brief
|
||||
* \author Han Hu
|
||||
* \modified by Jiarui Xu, Ze Liu
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cmath>
|
||||
#include <vector>
|
||||
|
||||
void similarity_compute_forward(
|
||||
const at::Tensor key,
|
||||
const at::Tensor query,
|
||||
const at::Tensor pos_weight,
|
||||
const int batch_size,
|
||||
const int key_channels,
|
||||
const int query_channels,
|
||||
const int height,
|
||||
const int width,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const at::Tensor scale,
|
||||
const at::Tensor no_define_value,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int in_height,
|
||||
const int in_width,
|
||||
const int sim_method,
|
||||
at::Tensor output,
|
||||
const int key_offset,
|
||||
const int query_offset);
|
||||
|
||||
void similarity_compute_backward(
|
||||
const at::Tensor key,
|
||||
const at::Tensor query,
|
||||
const at::Tensor output_grad,
|
||||
const int batch_size,
|
||||
const int key_channels,
|
||||
const int query_channels,
|
||||
const int height,
|
||||
const int width,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const int key_per_group,
|
||||
const at::Tensor scale,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int in_height,
|
||||
const int in_width,
|
||||
const int sim_method,
|
||||
at::Tensor key_grad,
|
||||
at::Tensor query_grad,
|
||||
const int key_grad_offset,
|
||||
const int query_grad_offset);
|
||||
|
||||
void aggregation_forward(
|
||||
const at::Tensor value,
|
||||
const at::Tensor softmax_data,
|
||||
const int batch_size,
|
||||
const int value_channels,
|
||||
const int height,
|
||||
const int width,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int in_height,
|
||||
const int in_width,
|
||||
at::Tensor output,
|
||||
const int value_offset,
|
||||
const int output_offset);
|
||||
|
||||
void aggregation_value_backward(
|
||||
const at::Tensor softmax_data,
|
||||
const at::Tensor output_grad,
|
||||
const int batch_size,
|
||||
const int value_channels,
|
||||
const int height,
|
||||
const int width,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int in_height,
|
||||
const int in_width,
|
||||
at::Tensor value_grad,
|
||||
const int output_grad_offset,
|
||||
const int value_grad_offset);
|
||||
|
||||
void aggregation_softmax_backward(
|
||||
const at::Tensor value,
|
||||
const at::Tensor output_grad,
|
||||
const int batch_size,
|
||||
const int value_channels,
|
||||
const int height,
|
||||
const int width,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int in_height,
|
||||
const int in_width,
|
||||
at::Tensor softmax_grad,
|
||||
const int value_offset,
|
||||
const int output_grad_offset);
|
||||
|
||||
|
||||
int local_relation_forward_cuda(
|
||||
at::Tensor query,
|
||||
at::Tensor key,
|
||||
at::Tensor value,
|
||||
at::Tensor pos_weight,
|
||||
at::Tensor scale,
|
||||
at::Tensor no_define_value,
|
||||
at::Tensor output,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int batch_step,
|
||||
const int norm_method,
|
||||
const int sim_method)
|
||||
{
|
||||
query = query.contiguous();
|
||||
key = key.contiguous();
|
||||
value = value.contiguous();
|
||||
pos_weight = pos_weight.contiguous();
|
||||
|
||||
const int query_channels = query.size(1);
|
||||
const int key_channels = key.size(1);
|
||||
const int value_channels = value.size(1);
|
||||
const int batch_size = key.size(0);
|
||||
const int height = query.size(2);
|
||||
const int width = query.size(3);
|
||||
const int in_height = key.size(2);
|
||||
const int in_width = key.size(3);
|
||||
|
||||
const int batch_step_ = std::min(batch_size, batch_step);
|
||||
const int sim_size = batch_step_ * num_group * kernel_height * kernel_width * height * width;
|
||||
|
||||
const int key_step = batch_step_ * key_channels * in_height * in_width;
|
||||
const int query_step = batch_step_ * query_channels * height * width;
|
||||
const int value_step = batch_step_ * value_channels * in_height * in_width;
|
||||
const int output_step = batch_step_ * value_channels * height * width;
|
||||
|
||||
at::Tensor sim_buffer = at::zeros({batch_step_ * num_group, kernel_height * kernel_width, height * width},
|
||||
query.options());
|
||||
|
||||
at::Tensor softmax_buffer = at::zeros({batch_step_ * num_group, kernel_height * kernel_width, height * width},
|
||||
query.options());
|
||||
|
||||
at::Tensor sum_softmax_buffer = at::zeros({batch_step_ * num_group, height * width});
|
||||
|
||||
int M = (batch_size - 1) / batch_step_ + 1;
|
||||
for (int i = 0; i < M; ++i) {
|
||||
int cur_batch_step = batch_step_;
|
||||
if (i == M - 1) {
|
||||
cur_batch_step = batch_size - (M - 1) * batch_step_;
|
||||
if (cur_batch_step != batch_step_) {
|
||||
sim_buffer = at::zeros({cur_batch_step * num_group, kernel_height * kernel_width, height * width}, query.options());
|
||||
softmax_buffer = at::zeros({cur_batch_step * num_group, kernel_height * kernel_width, height * width},query.options());
|
||||
sum_softmax_buffer = at::zeros({cur_batch_step * num_group, height * width}, query.options());
|
||||
}
|
||||
|
||||
// TORCH_CHECK(cur_batch_step % batch_step_ == 0, "batch_step must be divided by batch_size");
|
||||
}
|
||||
similarity_compute_forward(key, query, pos_weight, cur_batch_step,
|
||||
key_channels, query_channels, height, width,
|
||||
kernel_height, kernel_width, num_group, scale, no_define_value,
|
||||
dilate, stride, in_height, in_width, sim_method, sim_buffer,
|
||||
key_step * i, query_step * i);
|
||||
|
||||
// softmax
|
||||
if (norm_method == 0) {
|
||||
softmax_buffer = sim_buffer.softmax(1);
|
||||
}
|
||||
else {
|
||||
AT_ERROR("Not implemented yet");
|
||||
}
|
||||
|
||||
aggregation_forward(value, softmax_buffer, cur_batch_step,
|
||||
value_channels, height, width, kernel_height, kernel_width,
|
||||
num_group, dilate, stride, in_height, in_width, output, value_step * i, output_step * i);
|
||||
}
|
||||
|
||||
return 1;
|
||||
|
||||
}
|
||||
|
||||
int local_relation_backward_cuda(
|
||||
at::Tensor query,
|
||||
at::Tensor key,
|
||||
at::Tensor value,
|
||||
at::Tensor pos_weight,
|
||||
at::Tensor scale,
|
||||
at::Tensor no_define_value,
|
||||
at::Tensor output_grad,
|
||||
at::Tensor query_grad,
|
||||
at::Tensor key_grad,
|
||||
at::Tensor value_grad,
|
||||
at::Tensor pos_weight_grad,
|
||||
const int kernel_height,
|
||||
const int kernel_width,
|
||||
const int num_group,
|
||||
const int dilate,
|
||||
const int stride,
|
||||
const int batch_step,
|
||||
const int norm_method,
|
||||
const int sim_method)
|
||||
{
|
||||
query = query.contiguous();
|
||||
key = key.contiguous();
|
||||
value = value.contiguous();
|
||||
pos_weight = pos_weight.contiguous();
|
||||
|
||||
output_grad = output_grad.contiguous();
|
||||
query_grad = query_grad.contiguous();
|
||||
key_grad = key_grad.contiguous();
|
||||
value_grad = value_grad.contiguous();
|
||||
pos_weight_grad = pos_weight_grad.contiguous();
|
||||
|
||||
const int query_channels = query.size(1);
|
||||
const int key_channels = key.size(1);
|
||||
const int value_channels = value.size(1);
|
||||
const int batch_size = key.size(0);
|
||||
const int height = query.size(2);
|
||||
const int width = query.size(3);
|
||||
const int in_height = key.size(2);
|
||||
const int in_width = key.size(3);
|
||||
const int key_per_group = query_channels / num_group;
|
||||
|
||||
const int batch_step_ = std::min(batch_size, batch_step);
|
||||
const int sim_size = batch_step_ * num_group * kernel_height * kernel_width * height * width;
|
||||
|
||||
const int key_step = batch_step_ * key_channels * in_height * in_width;
|
||||
const int query_step = batch_step_ * query_channels * height * width;
|
||||
const int value_step = batch_step_ * value_channels * in_height * in_width;
|
||||
const int output_step = batch_step_ * value_channels * height * width;
|
||||
|
||||
at::Tensor sim_buffer = at::zeros({batch_step_ * num_group, kernel_height * kernel_width, height * width},
|
||||
query.options());
|
||||
|
||||
at::Tensor softmax_buffer = at::zeros({batch_step_ * num_group, kernel_height * kernel_width, height * width},
|
||||
query.options());
|
||||
|
||||
at::Tensor sum_softmax_buffer = at::zeros({batch_step_ * num_group, height * width},
|
||||
query.options());
|
||||
|
||||
at::Tensor sim_grad_buffer = at::zeros({batch_step_ * num_group, kernel_height * kernel_width, height * width},
|
||||
query.options());
|
||||
|
||||
int M = (batch_size - 1) / batch_step_ + 1;
|
||||
|
||||
const int pos_weight_size = num_group * kernel_height * kernel_width;
|
||||
|
||||
for (int i = 0; i < M; ++i) {
|
||||
int cur_batch_step = batch_step_;
|
||||
if (i == M - 1) {
|
||||
cur_batch_step = batch_size - (M - 1) * batch_step_;
|
||||
if (cur_batch_step != batch_step_) {
|
||||
sim_buffer = at::zeros({cur_batch_step * num_group, kernel_height * kernel_width, height * width}, query.options());
|
||||
softmax_buffer = at::zeros({cur_batch_step * num_group, kernel_height * kernel_width, height * width},query.options());
|
||||
sum_softmax_buffer = at::zeros({cur_batch_step * num_group, height * width}, query.options());
|
||||
sim_grad_buffer = at::zeros({cur_batch_step * num_group, kernel_height * kernel_width, height * width}, query.options());
|
||||
}
|
||||
// TORCH_CHECK(cur_batch_step % batch_step_ == 0, "batch_step must be divided by batch_size");
|
||||
}
|
||||
|
||||
similarity_compute_forward(key, query, pos_weight, cur_batch_step,
|
||||
key_channels, query_channels, height, width,
|
||||
kernel_height, kernel_width, num_group, scale, no_define_value,
|
||||
dilate, stride, in_height, in_width, sim_method, sim_buffer,
|
||||
key_step * i, query_step * i);
|
||||
|
||||
// softmax
|
||||
|
||||
if (norm_method == 0) {
|
||||
softmax_buffer = sim_buffer.softmax(1);
|
||||
}
|
||||
else {
|
||||
AT_ERROR("Not implemented yet");
|
||||
}
|
||||
|
||||
aggregation_value_backward(softmax_buffer, output_grad, cur_batch_step,
|
||||
value_channels, height, width, kernel_height, kernel_width,
|
||||
num_group, dilate, stride, in_height, in_width, value_grad,
|
||||
output_step * i, value_step * i);
|
||||
|
||||
aggregation_softmax_backward(value, output_grad, cur_batch_step,
|
||||
value_channels, height, width, kernel_height, kernel_width,
|
||||
num_group, dilate, stride, in_height, in_width, sim_buffer,
|
||||
value_step * i, output_step * i);
|
||||
|
||||
if (norm_method == 0) {
|
||||
sum_softmax_buffer = (softmax_buffer * sim_buffer).sum(1, true);
|
||||
sim_grad_buffer = softmax_buffer * (sim_buffer - sum_softmax_buffer);
|
||||
}
|
||||
else {
|
||||
AT_ERROR("Not implemented yet");
|
||||
}
|
||||
|
||||
similarity_compute_backward(key, query, sim_grad_buffer, cur_batch_step,
|
||||
key_channels, query_channels, height, width,
|
||||
kernel_height, kernel_width, num_group, key_per_group, scale,
|
||||
dilate, stride, in_height, in_width, sim_method, key_grad, query_grad,
|
||||
key_step * i, query_step * i);
|
||||
|
||||
pos_weight_grad += sim_grad_buffer.view({cur_batch_step, num_group, kernel_height, kernel_width, height * width}).sum(4).sum(0);
|
||||
|
||||
}
|
||||
|
||||
return 1;
|
||||
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("local_relation_forward_cuda", &local_relation_forward_cuda,
|
||||
"local relation forward (CUDA)");
|
||||
m.def("local_relation_backward_cuda", &local_relation_backward_cuda,
|
||||
"local relation backward (CUDA)");
|
||||
}
|
||||
1004
ops/local_relation/src/local_relation_cuda_kernel.cu
Normal file
1004
ops/local_relation/src/local_relation_cuda_kernel.cu
Normal file
File diff suppressed because it is too large
Load Diff
Loading…
Reference in New Issue
Block a user