[
{
"path": ".gitignore",
"content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n"
},
{
"path": "LICENSE",
"content": "MIT License\n\nCopyright (c) 2022 Megvii Inc.\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n"
},
{
"path": "README.md",
"content": "## BSRT: Improving Burst Super-Resolution with Swin Transformer and Flow-Guided Deformable Alignment (CVPRW 2022)\n[](https://paperswithcode.com/sota/burst-image-super-resolution-on-burstsr?p=bsrt-improving-burst-super-resolution-with) [](https://paperswithcode.com/sota/burst-image-super-resolution-on?p=bsrt-improving-burst-super-resolution-with) \n\n#### [BSRT](https://arxiv.org/abs/2204.08332), the winner of the NTIRE 2022 Burst Super-Resolution Challenge Real-World Track. \nYou can also find our winner method in NTIRE 2021 Burst Super-Resolution Challenge [here](https://github.com/Algolzw/EBSR).\n\n> This work addresses the Burst Super-Resolution (BurstSR) task using a new architecture, which requires restoring a high-quality image from a sequence of noisy, misaligned, and low-resolution RAW bursts. To overcome the challenges in BurstSR, we propose a **B**urst **S**uper-**R**esolution **T**ransformer (**BSRT**), which can significantly improve the capability of extracting inter-frame information and reconstruction. To achieve this goal, we propose a Pyramid Flow-Guided Deformable Convolution Network (Pyramid FG-DCN) and incorporate Swin Transformer Blocks and Groups as our main backbone. More specifically, we combine optical flows and deformable convolutions, hence our BSRT can handle misalignment and aggregate the potential texture information in multi-frames more efficiently. In addition, our Transformer-based structure can capture long-range dependency to further improve the performance. The evaluation on both synthetic and real-world tracks demonstrates that our approach achieves a new state-of-the-art in BurstSR task. Further, our BSRT wins the championship in the NTIRE2022 Burst Super-Resolution Challenge.\n\n\n#### Comparison with State-of-the-art Burst Super-Resolution Methods\n\n\n\n\n\n## Overview Architecture\n\n\n\n## Dependencies\n- OS: Ubuntu 18.04\n- Python: Python 3.7\n- nvidia :\n - cuda: 10.1\n - cudnn: 7.6.1\n- Other reference requirements\n\n## Quick Start\n1.Create a conda virtual environment and activate it\n```python3\nconda create -n pytorch_1.6 python=3.7\nsource activate pytorch_1.6\n```\n2.Install PyTorch and torchvision following the official instructions\n```python3\nconda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch\n```\n3.Install build requirements\n```python3\npip3 install -r requirements.txt\n```\n4.Install DCN\n```python3\ncd DCNv2\npython3 setup.py build develop # build\npython3 test.py # run examples and check\n```\n\n## Training\n\nWe provide all pretrained model weights [here](https://drive.google.com/file/d/1Bv1ZwoE3s8trhG--wjB0Yt6WJIQPpvsn/view?usp=sharing). \n\n#### For Synthetic data\n\n```python3\ncd code/synthetic/bsrt\n# Modify the root path of training dataset and model etc.\n# The number of GPUs should be more than 1\npython main.py --n_GPUs 8 --print_every 40 --lr 0.0001 --decay 150-300 --save bsrt_tiny --model BSRT --fp16 --model_level S --swinfeature --batch_size 32 --burst_size 14 --patch_size 256\n```\n\n#### For Real-World data\n\n```python3\ncd code/real/bsrt\n# Modify the root path of training dataset and model etc.\n# The number of GPUs should be more than 1\npython main.py --n_GPUs 8 --print_every 20 --lr 0.00005 --decay 40-80 --save bsrt_tiny --model BSRT --fp16 --model_level S --swinfeature --batch_size 8 --burst_size 14 --patch_size 80 --pre_train ../../synthetic/train_log/bsrt/real_models/bsrt_tiny/bsrt_best_epoch.pth \n```\n\nThe pretrained PWC-Net model can be downloaded [here](https://drive.google.com/file/d/1dD6vB9QN3qwmOBi3AGKzJbbSojwDDlgV/view?usp=sharing). \n\n## Test\n\n#### For Synthetic data\n```python3\n# Modify the path of test dataset and the path of the trained model\npython test_synburst.py --n_GPUs 1 --model BSRT --model_level S --swinfeature --burst_size 14 --patch_size 384 --pre_train ../train_log/bsrt/real_models/bsrt_tiny/bsrt_best_epoch.pth --root /data/dataset/ntire21/burstsr/synthetic\n```\n\n#### For Real-World data\n```python3\n# Modify the path of test dataset and the path of the trained model\npython test_real.py --n_GPUs 1 --model BSRT --model_level S --swinfeature --batch_size 1 --burst_size 14 --patch_size 80 --pre_train ../train_log/bsrt/real_models/bsrt_tiny/bsrtbest_epoch.pth --root /data/dataset/ntire21/burstsr/real\n```\n\n## Results\n\n### Comparison on Synthetic dataset\n\n\n### Comparison on Real-World dataset\n\n\n\n## Citations\nIf our code helps your research or work, please consider citing our paper.\nThe following is a BibTeX reference.\n\n```\n@inproceedings{luo2022bsrt,\n title={BSRT: Improving Burst Super-Resolution with Swin Transformer and Flow-Guided Deformable Alignment},\n author={Luo, Ziwei and Li, Youwei and Cheng, Shen and Yu, Lei and Wu, Qi and Wen, Zhihong and Fan, Haoqiang and Sun, Jian and Liu, Shuaicheng},\n booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition},\n pages={998--1008},\n year={2022}\n}\n```\n\n## Contact\nemail: [ziwei.ro@gmail.com]\n"
},
{
"path": "code/real/bsrt/README.md",
"content": "# BSRT: Improving Burst Super-Resolution with Swin Transformer and Flow-Guided Deformable Alignment (Real-World)\n\n## Dependencies\n- OS: Ubuntu 18.04\n- Python: Python 3.7\n- nvidia :\n - cuda: 10.1\n - cudnn: 7.6.1\n- Other reference requirements\n\n## Quick Start\n1.Create a conda virtual environment and activate it\n```python3\nconda create -n pytorch_1.6 python=3.7\nsource activate pytorch_1.6\n```\n2.Install PyTorch and torchvision following the official instructions\n```python3\nconda install pytorch==1.6.0 torchvision==0.7.0 cudatoolkit=10.1 -c pytorch\n```\n3.Install build requirements\n```python3\npip3 install -r requirements.txt\n```\n4.Install DCN\n```python3\ncd DCNv2\npython3 setup.py build develop # build\npython3 test.py # run examples and check\n```\n\n## Training\n\nThe pretrained PWC-Net model can be downloaded [here](https://drive.google.com/file/d/1dD6vB9QN3qwmOBi3AGKzJbbSojwDDlgV/view?usp=sharing). \n\n```python3\n# Modify the root path of training dataset and model etc.\n# The number of GPUs should be more than 1\npython main.py --n_GPUs 8 --print_every 20 --lr 0.00004 --decay 40-80 --save bsrt_tiny --model BSRT --fp16 --model_level S --swinfeature --batch_size 8 --burst_size 14 --patch_size 80 --pre_train ../../synthetic/train_log/bsrt/real_models/bsrt_tiny/bsrt_best_epoch.pth \n```\n## Test\n```python3\n# Modify the path of test dataset and the path of the trained model\npython test_real.py --n_GPUs 1 --model BSRT --model_level S --swinfeature --batch_size 1 --burst_size 14 --patch_size 80 --pre_train ../train_log/bsrt/real_models/bsrt_tiny/bsrtbest_epoch.pth --root /data/dataset/ntire21/burstsr/real\n```"
},
{
"path": "code/real/bsrt/data_processing/__init__.py",
"content": ""
},
{
"path": "code/real/bsrt/data_processing/camera_pipeline.py",
"content": "import torch\nimport random\nimport math\nimport cv2 as cv\nimport numpy as np\nimport utils.data_format_utils as df_utils\n\"\"\" Based on http://timothybrooks.com/tech/unprocessing\nFunctions for forward and inverse camera pipeline. All functions input a torch float tensor of shape (c, h, w).\nAdditionally, some also support batch operations, i.e. inputs of shape (b, c, h, w)\n\"\"\"\n\n\ndef random_ccm():\n \"\"\"Generates random RGB -> Camera color correction matrices.\"\"\"\n # Takes a random convex combination of XYZ -> Camera CCMs.\n xyz2cams = [[[1.0234, -0.2969, -0.2266],\n [-0.5625, 1.6328, -0.0469],\n [-0.0703, 0.2188, 0.6406]],\n [[0.4913, -0.0541, -0.0202],\n [-0.613, 1.3513, 0.2906],\n [-0.1564, 0.2151, 0.7183]],\n [[0.838, -0.263, -0.0639],\n [-0.2887, 1.0725, 0.2496],\n [-0.0627, 0.1427, 0.5438]],\n [[0.6596, -0.2079, -0.0562],\n [-0.4782, 1.3016, 0.1933],\n [-0.097, 0.1581, 0.5181]]]\n\n num_ccms = len(xyz2cams)\n xyz2cams = torch.tensor(xyz2cams)\n\n weights = torch.FloatTensor(num_ccms, 1, 1).uniform_(0.0, 1.0)\n weights_sum = weights.sum()\n xyz2cam = (xyz2cams * weights).sum(dim=0) / weights_sum\n\n # Multiplies with RGB -> XYZ to get RGB -> Camera CCM.\n rgb2xyz = torch.tensor([[0.4124564, 0.3575761, 0.1804375],\n [0.2126729, 0.7151522, 0.0721750],\n [0.0193339, 0.1191920, 0.9503041]])\n rgb2cam = torch.mm(xyz2cam, rgb2xyz)\n\n # Normalizes each row.\n rgb2cam = rgb2cam / rgb2cam.sum(dim=-1, keepdims=True)\n return rgb2cam\n\n\ndef random_gains():\n \"\"\"Generates random gains for brightening and white balance.\"\"\"\n # RGB gain represents brightening.\n rgb_gain = 1.0 / random.gauss(mu=0.8, sigma=0.1)\n\n # Red and blue gains represent white balance.\n red_gain = random.uniform(1.9, 2.4)\n blue_gain = random.uniform(1.5, 1.9)\n return rgb_gain, red_gain, blue_gain\n\n\ndef apply_smoothstep(image):\n \"\"\"Apply global tone mapping curve.\"\"\"\n image_out = 3 * image**2 - 2 * image**3\n return image_out\n\n\ndef invert_smoothstep(image):\n \"\"\"Approximately inverts a global tone mapping curve.\"\"\"\n image = image.clamp(0.0, 1.0)\n return 0.5 - torch.sin(torch.asin(1.0 - 2.0 * image) / 3.0)\n\n\ndef gamma_expansion(image):\n \"\"\"Converts from gamma to linear space.\"\"\"\n # Clamps to prevent numerical instability of gradients near zero.\n return image.clamp(1e-8) ** 2.2\n\n\ndef gamma_compression(image):\n \"\"\"Converts from linear to gammaspace.\"\"\"\n # Clamps to prevent numerical instability of gradients near zero.\n return image.clamp(1e-8) ** (1.0 / 2.2)\n\n\ndef apply_ccm(image, ccm):\n \"\"\"Applies a color correction matrix.\"\"\"\n assert image.dim() == 3 and image.shape[0] == 3\n\n shape = image.shape\n image = image.view(3, -1)\n ccm = ccm.to(image.device).type_as(image)\n\n image = torch.mm(ccm, image)\n\n return image.view(shape)\n\n\ndef apply_gains(image, rgb_gain, red_gain, blue_gain):\n \"\"\"Inverts gains while safely handling saturated pixels.\"\"\"\n assert image.dim() == 3 and image.shape[0] in [3, 4]\n\n if image.shape[0] == 3:\n gains = torch.tensor([red_gain, 1.0, blue_gain]) * rgb_gain\n else:\n gains = torch.tensor([red_gain, 1.0, 1.0, blue_gain]) * rgb_gain\n gains = gains.view(-1, 1, 1)\n gains = gains.to(image.device).type_as(image)\n\n return (image * gains).clamp(0.0, 1.0)\n\n\ndef safe_invert_gains(image, rgb_gain, red_gain, blue_gain):\n \"\"\"Inverts gains while safely handling saturated pixels.\"\"\"\n assert image.dim() == 3 and image.shape[0] == 3\n\n gains = torch.tensor([1.0 / red_gain, 1.0, 1.0 / blue_gain]) / rgb_gain\n gains = gains.view(-1, 1, 1)\n\n # Prevents dimming of saturated pixels by smoothly masking gains near white.\n gray = image.mean(dim=0, keepdims=True)\n inflection = 0.9\n mask = ((gray - inflection).clamp(0.0) / (1.0 - inflection)) ** 2.0\n\n safe_gains = torch.max(mask + (1.0 - mask) * gains, gains)\n return image * safe_gains\n\n\ndef mosaic(image, mode='rggb'):\n \"\"\"Extracts RGGB Bayer planes from an RGB image.\"\"\"\n shape = image.shape\n if image.dim() == 3:\n image = image.unsqueeze(0)\n\n if mode == 'rggb':\n red = image[:, 0, 0::2, 0::2]\n green_red = image[:, 1, 0::2, 1::2]\n green_blue = image[:, 1, 1::2, 0::2]\n blue = image[:, 2, 1::2, 1::2]\n image = torch.stack((red, green_red, green_blue, blue), dim=1)\n elif mode == 'grbg':\n green_red = image[:, 1, 0::2, 0::2]\n red = image[:, 0, 0::2, 1::2]\n blue = image[:, 2, 0::2, 1::2]\n green_blue = image[:, 1, 1::2, 1::2]\n\n image = torch.stack((green_red, red, blue, green_blue), dim=1)\n\n if len(shape) == 3:\n return image.view((4, shape[-2] // 2, shape[-1] // 2))\n else:\n return image.view((-1, 4, shape[-2] // 2, shape[-1] // 2))\n\n\ndef demosaic(image):\n assert isinstance(image, torch.Tensor)\n image = image.clamp(0.0, 1.0) * 255\n\n if image.dim() == 4:\n num_images = image.dim()\n batch_input = True\n else:\n num_images = 1\n batch_input = False\n image = image.unsqueeze(0)\n\n # Generate single channel input for opencv\n im_sc = torch.zeros((num_images, image.shape[-2] * 2, image.shape[-1] * 2, 1))\n im_sc[:, ::2, ::2, 0] = image[:, 0, :, :]\n im_sc[:, ::2, 1::2, 0] = image[:, 1, :, :]\n im_sc[:, 1::2, ::2, 0] = image[:, 2, :, :]\n im_sc[:, 1::2, 1::2, 0] = image[:, 3, :, :]\n\n im_sc = im_sc.numpy().astype(np.uint8)\n\n out = []\n\n for im in im_sc:\n # cv.imwrite('frames/tmp.png', im)\n im_dem_np = cv.cvtColor(im, cv.COLOR_BAYER_BG2RGB)#_VNG)\n\n # Convert to torch image\n im_t = df_utils.npimage_to_torch(im_dem_np, input_bgr=False)\n out.append(im_t)\n\n if batch_input:\n return torch.stack(out, dim=0)\n else:\n return out[0]\n\n\ndef random_noise_levels():\n \"\"\"Generates random noise levels from a log-log linear distribution.\"\"\"\n log_min_shot_noise = math.log(0.0001)\n log_max_shot_noise = math.log(0.012)\n log_shot_noise = random.uniform(log_min_shot_noise, log_max_shot_noise)\n shot_noise = math.exp(log_shot_noise)\n\n line = lambda x: 2.18 * x + 1.20\n log_read_noise = line(log_shot_noise) + random.gauss(mu=0.0, sigma=0.26)\n read_noise = math.exp(log_read_noise)\n return shot_noise, read_noise\n\n\ndef add_noise(image, shot_noise=0.01, read_noise=0.0005):\n \"\"\"Adds random shot (proportional to image) and read (independent) noise.\"\"\"\n variance = image * shot_noise + read_noise\n noise = torch.FloatTensor(image.shape).normal_().to(image.device)*variance.sqrt()\n return image + noise\n\n\ndef process_linear_image_rgb(image, meta_info, return_np=False):\n image = apply_gains(image, meta_info['rgb_gain'], meta_info['red_gain'], meta_info['blue_gain'])\n image = apply_ccm(image, meta_info['cam2rgb'])\n\n if meta_info['gamma']:\n image = gamma_compression(image)\n\n if meta_info['smoothstep']:\n image = apply_smoothstep(image)\n\n image = image.clamp(0.0, 1.0)\n\n if return_np:\n image = df_utils.torch_to_npimage(image)\n return image\n\n\ndef process_linear_image_raw(image, meta_info):\n image = apply_gains(image, meta_info['rgb_gain'], meta_info['red_gain'], meta_info['blue_gain'])\n image = demosaic(image)\n image = apply_ccm(image, meta_info['cam2rgb'])\n\n if meta_info['gamma']:\n image = gamma_compression(image)\n\n if meta_info['smoothstep']:\n image = apply_smoothstep(image)\n return image.clamp(0.0, 1.0)\n"
},
{
"path": "code/real/bsrt/data_processing/synthetic_burst_generation.py",
"content": "import torch\nimport random\nimport cv2\nimport numpy as np\nimport torch.nn.functional as F\nfrom data_processing.camera_pipeline import *\nfrom utils.data_format_utils import torch_to_numpy, numpy_to_torch\n\n\ndef random_crop(frames, crop_sz):\n \"\"\" Extract a random crop of size crop_sz from the input frames. If the crop_sz is larger than the input image size,\n then the largest possible crop of same aspect ratio as crop_sz will be extracted from frames, and upsampled to\n crop_sz.\n \"\"\"\n if not isinstance(crop_sz, (tuple, list)):\n crop_sz = (crop_sz, crop_sz)\n crop_sz = torch.tensor(crop_sz).float()\n\n shape = frames.shape\n\n # Select scale_factor. Ensure the crop fits inside the image\n max_scale_factor = torch.tensor(shape[-2:]).float() / crop_sz\n max_scale_factor = max_scale_factor.min().item()\n\n if max_scale_factor < 1.0:\n scale_factor = max_scale_factor\n else:\n scale_factor = 1.0\n\n # Extract the crop\n orig_crop_sz = (crop_sz * scale_factor).floor()\n\n assert orig_crop_sz[-2] <= shape[-2] and orig_crop_sz[-1] <= shape[-1], 'Bug in crop size estimation!'\n\n r1 = random.randint(0, shape[-2] - orig_crop_sz[-2])\n c1 = random.randint(0, shape[-1] - orig_crop_sz[-1])\n\n r2 = r1 + orig_crop_sz[0].int().item()\n c2 = c1 + orig_crop_sz[1].int().item()\n\n frames_crop = frames[:, r1:r2, c1:c2]\n\n # Resize to crop_sz\n if scale_factor < 1.0:\n frames_crop = F.interpolate(frames_crop.unsqueeze(0), size=crop_sz.int().tolist(), mode='bilinear', align_corners=False).squeeze(0)\n return frames_crop\n\n\ndef rgb2rawburst(image, burst_size, downsample_factor=1, burst_transformation_params=None,\n image_processing_params=None, interpolation_type='bilinear'):\n \"\"\" Generates a synthetic LR RAW burst from the input image. The input sRGB image is first converted to linear\n sensor space using an inverse camera pipeline. A LR burst is then generated by applying random\n transformations defined by burst_transformation_params to the input image, and downsampling it by the\n downsample_factor. The generated burst is then mosaicekd and corrputed by random noise.\n \"\"\"\n\n if image_processing_params is None:\n image_processing_params = {}\n\n _defaults = {'random_ccm': True, 'random_gains': True, 'smoothstep': True, 'gamma': True, 'add_noise': True}\n for k, v in _defaults.items():\n if k not in image_processing_params:\n image_processing_params[k] = v\n\n # Sample camera pipeline params\n if image_processing_params['random_ccm']:\n rgb2cam = random_ccm()\n else:\n rgb2cam = torch.eye(3).float()\n cam2rgb = rgb2cam.inverse()\n\n # Sample gains\n if image_processing_params['random_gains']:\n rgb_gain, red_gain, blue_gain = random_gains()\n else:\n rgb_gain, red_gain, blue_gain = (1.0, 1.0, 1.0)\n\n # Approximately inverts global tone mapping.\n use_smoothstep = image_processing_params['smoothstep']\n if use_smoothstep:\n image = invert_smoothstep(image)\n\n # Inverts gamma compression.\n use_gamma = image_processing_params['gamma']\n if use_gamma:\n image = gamma_expansion(image)\n\n # Inverts color correction.\n image = apply_ccm(image, rgb2cam)\n\n # Approximately inverts white balance and brightening.\n image = safe_invert_gains(image, rgb_gain, red_gain, blue_gain)\n\n # Clip saturated pixels.\n image = image.clamp(0.0, 1.0)\n\n # Generate LR burst\n image_burst_rgb, flow_vectors = single2lrburst(image, burst_size=burst_size,\n downsample_factor=downsample_factor,\n transformation_params=burst_transformation_params,\n interpolation_type=interpolation_type)\n\n # mosaic\n image_burst = mosaic(image_burst_rgb.clone())\n\n # Add noise\n if image_processing_params['add_noise']:\n shot_noise_level, read_noise_level = random_noise_levels()\n image_burst = add_noise(image_burst, shot_noise_level, read_noise_level)\n else:\n shot_noise_level = 0\n read_noise_level = 0\n\n # Clip saturated pixels.\n image_burst = image_burst.clamp(0.0, 1.0)\n\n meta_info = {'rgb2cam': rgb2cam, 'cam2rgb': cam2rgb, 'rgb_gain': rgb_gain, 'red_gain': red_gain,\n 'blue_gain': blue_gain, 'smoothstep': use_smoothstep, 'gamma': use_gamma,\n 'shot_noise_level': shot_noise_level, 'read_noise_level': read_noise_level}\n return image_burst, image, image_burst_rgb, flow_vectors, meta_info\n\n\ndef get_tmat(image_shape, translation, theta, shear_values, scale_factors):\n \"\"\" Generates a transformation matrix corresponding to the input transformation parameters \"\"\"\n im_h, im_w = image_shape\n\n t_mat = np.identity(3)\n\n t_mat[0, 2] = translation[0]\n t_mat[1, 2] = translation[1]\n t_rot = cv2.getRotationMatrix2D((im_w * 0.5, im_h * 0.5), theta, 1.0)\n t_rot = np.concatenate((t_rot, np.array([0.0, 0.0, 1.0]).reshape(1, 3)))\n\n t_shear = np.array([[1.0, shear_values[0], -shear_values[0] * 0.5 * im_w],\n [shear_values[1], 1.0, -shear_values[1] * 0.5 * im_h],\n [0.0, 0.0, 1.0]])\n\n t_scale = np.array([[scale_factors[0], 0.0, 0.0],\n [0.0, scale_factors[1], 0.0],\n [0.0, 0.0, 1.0]])\n\n t_mat = t_scale @ t_rot @ t_shear @ t_mat\n\n t_mat = t_mat[:2, :]\n\n return t_mat\n\n\ndef single2lrburst(image, burst_size, downsample_factor=1, transformation_params=None,\n interpolation_type='bilinear'):\n \"\"\" Generates a burst of size burst_size from the input image by applying random transformations defined by\n transformation_params, and downsampling the resulting burst by downsample_factor.\n \"\"\"\n\n if interpolation_type == 'bilinear':\n interpolation = cv2.INTER_LINEAR\n elif interpolation_type == 'lanczos':\n interpolation = cv2.INTER_LANCZOS4\n else:\n raise ValueError\n\n normalize = False\n if isinstance(image, torch.Tensor):\n if image.max() < 2.0:\n image = image * 255.0\n normalize = True\n image = torch_to_numpy(image).astype(np.uint8)\n\n burst = []\n sample_pos_inv_all = []\n\n rvs, cvs = torch.meshgrid([torch.arange(0, image.shape[0]),\n torch.arange(0, image.shape[1])])\n\n sample_grid = torch.stack((cvs, rvs, torch.ones_like(cvs)), dim=-1).float()\n\n for i in range(burst_size):\n if i == 0:\n # For base image, do not apply any random transformations. We only translate the image to center the\n # sampling grid\n shift = (downsample_factor / 2.0) - 0.5\n translation = (shift, shift)\n theta = 0.0\n shear_factor = (0.0, 0.0)\n scale_factor = (1.0, 1.0)\n else:\n # Sample random image transformation parameters\n max_translation = transformation_params.get('max_translation', 0.0)\n\n if max_translation <= 0.01:\n shift = (downsample_factor / 2.0) - 0.5\n translation = (shift, shift)\n else:\n translation = (random.uniform(-max_translation, max_translation),\n random.uniform(-max_translation, max_translation))\n\n max_rotation = transformation_params.get('max_rotation', 0.0)\n theta = random.uniform(-max_rotation, max_rotation)\n\n max_shear = transformation_params.get('max_shear', 0.0)\n shear_x = random.uniform(-max_shear, max_shear)\n shear_y = random.uniform(-max_shear, max_shear)\n shear_factor = (shear_x, shear_y)\n\n max_ar_factor = transformation_params.get('max_ar_factor', 0.0)\n ar_factor = np.exp(random.uniform(-max_ar_factor, max_ar_factor))\n\n max_scale = transformation_params.get('max_scale', 0.0)\n scale_factor = np.exp(random.uniform(-max_scale, max_scale))\n\n scale_factor = (scale_factor, scale_factor * ar_factor)\n\n output_sz = (image.shape[1], image.shape[0])\n\n # Generate a affine transformation matrix corresponding to the sampled parameters\n t_mat = get_tmat((image.shape[0], image.shape[1]), translation, theta, shear_factor, scale_factor)\n t_mat_tensor = torch.from_numpy(t_mat)\n\n # Apply the sampled affine transformation\n image_t = cv2.warpAffine(image, t_mat, output_sz, flags=interpolation,\n borderMode=cv2.BORDER_CONSTANT)\n\n t_mat_tensor_3x3 = torch.cat((t_mat_tensor.float(), torch.tensor([0.0, 0.0, 1.0]).view(1, 3)), dim=0)\n t_mat_tensor_inverse = t_mat_tensor_3x3.inverse()[:2, :].contiguous()\n\n sample_pos_inv = torch.mm(sample_grid.view(-1, 3), t_mat_tensor_inverse.t().float()).view(\n *sample_grid.shape[:2], -1)\n\n if transformation_params.get('border_crop') is not None:\n border_crop = transformation_params.get('border_crop')\n\n image_t = image_t[border_crop:-border_crop, border_crop:-border_crop, :]\n sample_pos_inv = sample_pos_inv[border_crop:-border_crop, border_crop:-border_crop, :]\n\n # Downsample the image\n image_t = cv2.resize(image_t, None, fx=1.0 / downsample_factor, fy=1.0 / downsample_factor,\n interpolation=interpolation)\n sample_pos_inv = cv2.resize(sample_pos_inv.numpy(), None, fx=1.0 / downsample_factor,\n fy=1.0 / downsample_factor,\n interpolation=interpolation)\n\n sample_pos_inv = torch.from_numpy(sample_pos_inv).permute(2, 0, 1).contiguous()\n\n if normalize:\n image_t = numpy_to_torch(image_t).float() / 255.0\n else:\n image_t = numpy_to_torch(image_t).float()\n burst.append(image_t)\n sample_pos_inv_all.append(sample_pos_inv / downsample_factor)\n\n burst_images = torch.stack(burst)\n sample_pos_inv_all = torch.stack(sample_pos_inv_all)\n\n # Compute the flow vectors to go from the i'th burst image to the base image\n flow_vectors = sample_pos_inv_all - sample_pos_inv_all[:, :1, ...]\n\n return burst_images, flow_vectors\n"
},
{
"path": "code/real/bsrt/datasets/__init__.py",
"content": ""
},
{
"path": "code/real/bsrt/datasets/burstsr_dataset.py",
"content": "import os\nimport torch\nimport cv2\nimport numpy as np\nimport pickle as pkl\nimport torch.nn.functional as F\nimport random\nimport time\n\nclass SamsungRAWImage:\n @staticmethod\n def load(path):\n im_raw = cv2.imread('{}/im_raw.png'.format(path), cv2.IMREAD_UNCHANGED)\n\n im_raw = np.transpose(im_raw, (2, 0, 1)).astype(np.int16)\n im_raw = torch.from_numpy(im_raw)\n\n meta_data = pkl.load(open('{}/meta_info.pkl'.format(path), \"rb\", -1))\n\n return SamsungRAWImage(im_raw, meta_data['black_level'], meta_data['cam_wb'],\n meta_data['daylight_wb'], meta_data['color_matrix'], meta_data['exif_data'],\n meta_data.get('crop_info', None), meta_data.get('im_preview', None))\n\n def __init__(self, im_raw, black_level, cam_wb, daylight_wb, color_matrix, exif_data, crop_info=None,\n im_preview=None):\n self.im_raw = im_raw\n\n self.black_level = black_level\n self.cam_wb = cam_wb\n self.daylight_wb = daylight_wb\n self.color_matrix = color_matrix\n self.exif_data = exif_data\n self.crop_info = crop_info\n self.im_preview = im_preview\n\n self.norm_factor = 1023.0\n\n def get_all_meta_data(self):\n return {'black_level': self.black_level, 'cam_wb': self.cam_wb, 'daylight_wb': self.daylight_wb,\n 'color_matrix': self.color_matrix.tolist()}\n\n def get_exposure_time(self):\n return self.exif_data['Image ExposureTime'].values[0].decimal()\n\n def get_noise_profile(self):\n noise = self.exif_data['Image Tag 0xC761'].values\n noise = [n[0] for n in noise]\n noise = np.array(noise).reshape(3, 2)\n return noise\n\n def get_f_number(self):\n return self.exif_data['Image FNumber'].values[0].decimal()\n\n def get_iso(self):\n return self.exif_data['Image ISOSpeedRatings'].values[0]\n\n def get_image_data(self, substract_black_level=False, white_balance=False, normalize=False):\n im_raw = self.im_raw.float()\n\n if substract_black_level:\n im_raw = im_raw - torch.tensor(self.black_level).view(4, 1, 1)\n\n if white_balance:\n im_raw = im_raw * torch.tensor(self.cam_wb).view(4, 1, 1)\n\n if normalize:\n im_raw = im_raw / self.norm_factor\n\n\n return im_raw\n\n def shape(self):\n shape = (4, self.im_raw.shape[1], self.im_raw.shape[2])\n return shape\n\n def crop_image(self, r1, r2, c1, c2):\n self.im_raw = self.im_raw[:, r1:r2, c1:c2]\n\n def get_crop(self, r1, r2, c1, c2):\n im_raw = self.im_raw[:, r1:r2, c1:c2]\n\n if self.im_preview is not None:\n im_preview = self.im_preview[2*r1:2*r2, 2*c1:2*c2]\n else:\n im_preview = None\n\n return SamsungRAWImage(im_raw, self.black_level, self.cam_wb, self.daylight_wb, self.color_matrix,\n self.exif_data, im_preview=im_preview)\n\n def postprocess(self, return_np=True, norm_factor=None):\n # Convert to rgb\n # im = torch.from_numpy(self.im_raw.astype(np.float32))\n im = self.im_raw\n\n im = (im - torch.tensor(self.black_level).view(4, 1, 1)) * torch.tensor(self.cam_wb).view(4, 1, 1)\n\n if norm_factor is None:\n im = im / im.max()\n else:\n im = im / norm_factor\n\n im = torch.stack((im[0], (im[1] + im[2])/2, im[3]), dim=0)\n # im = torch.stack((im[0], im[1], im[3]), dim=0)\n\n im_out = im.clamp(0.0, 1.0)\n\n if return_np:\n im_out = im_out.permute(1, 2, 0).numpy() * 255.0\n im_out = im_out.astype(np.uint8)\n return im_out\n\n\nclass CanonImage:\n @staticmethod\n def load(path, split='train'):\n im_raw = cv2.imread('{}/im_raw.png'.format(path), cv2.IMREAD_UNCHANGED)\n im_raw = np.transpose(im_raw, (2, 0, 1)).astype(np.int16)\n im_raw = torch.from_numpy(im_raw)\n meta_data = pkl.load(open('{}/meta_info.pkl'.format(path), \"rb\", -1))\n\n return CanonImage(im_raw.float(), meta_data['black_level'], meta_data['cam_wb'],\n meta_data['daylight_wb'], meta_data['rgb_xyz_matrix'], meta_data.get('exif_data', None),\n meta_data.get('crop_info', None))\n\n def __init__(self, im_raw, black_level, cam_wb, daylight_wb, rgb_xyz_matrix, exif_data, crop_info=None):\n super(CanonImage, self).__init__()\n self.im_raw = im_raw\n\n if len(black_level) == 4:\n black_level = [black_level[0], black_level[1], black_level[3]]\n self.black_level = black_level\n\n if len(cam_wb) == 4:\n cam_wb = [cam_wb[0], cam_wb[1], cam_wb[3]]\n self.cam_wb = cam_wb\n\n if len(daylight_wb) == 4:\n daylight_wb = [daylight_wb[0], daylight_wb[1], daylight_wb[3]]\n self.daylight_wb = daylight_wb\n\n self.rgb_xyz_matrix = rgb_xyz_matrix\n self.xyz_srgb_matrix = torch.tensor([3.2404542, -1.5371385, -0.4985314,\n -0.9692660, 1.8760108, 0.0415560,\n 0.0556434, -0.2040259, 1.0572252]).view(3, 3)\n self.exif_data = exif_data\n self.crop_info = crop_info\n\n self.norm_factor = 16383\n\n def shape(self):\n shape = (3, self.im_raw.shape[1], self.im_raw.shape[2])\n return shape\n\n def get_all_meta_data(self):\n return {'black_level': self.black_level, 'cam_wb': self.cam_wb, 'daylight_wb': self.daylight_wb,\n 'rgb_xyz_matrix': self.rgb_xyz_matrix.tolist(), 'crop_info': self.crop_info,\n 'norm_factor': self.norm_factor}\n\n def get_exposure_time(self):\n return self.exif_data['EXIF ExposureTime'].values[0].decimal()\n\n def get_f_number(self):\n return self.exif_data['EXIF FNumber'].values[0].decimal()\n\n def get_iso(self):\n return self.exif_data['EXIF ISOSpeedRatings'].values[0]\n\n def get_image_data(self, substract_black_level=False, white_balance=False, normalize=False):\n im_raw = self.im_raw.float()\n\n if substract_black_level:\n im_raw = im_raw - torch.tensor(self.black_level).view(3, 1, 1)\n\n if white_balance:\n im_raw = im_raw * torch.tensor(self.cam_wb).view(3, 1, 1) / 1024.0\n\n if normalize:\n im_raw = im_raw / self.norm_factor\n\n return im_raw\n\n def set_image_data(self, im_data):\n self.im_raw = im_data\n\n def crop_image(self, r1, r2, c1, c2):\n self.im_raw = self.im_raw[:, r1:r2, c1:c2]\n\n def get_crop(self, r1, r2, c1, c2):\n im_raw = self.im_raw[:, r1:r2, c1:c2]\n return CanonImage(im_raw, self.black_level, self.cam_wb, self.daylight_wb, self.rgb_xyz_matrix,\n self.exif_data, self.crop_info)\n\n def set_crop_info(self, crop_info):\n self.crop_info = crop_info\n\n def resize(self, size=None, scale_factor=None):\n\n self.im_raw = F.interpolate(self.im_raw.unsqueeze(0), size=size, scale_factor=scale_factor,\n mode='bilinear').squeeze(0)\n\n def postprocess(self, return_np=True):\n # Convert to rgb\n im = self.im_raw\n\n im = (im - torch.tensor(self.black_level).view(3, 1, 1)).float() * torch.tensor(self.cam_wb).view(3, 1, 1)\n\n im_out = im / im.max()\n im_out = im_out.clamp(0.0, 1.0)\n\n if return_np:\n im_out = im_out.permute(1, 2, 0).numpy() * 255.0\n im_out = im_out.astype(np.uint8)\n return im_out\n\n\ndef load_txt(path):\n with open(path, 'r') as fh:\n out = [d.rstrip() for d in fh.readlines()]\n\n return out\n\n\nclass BurstSRDataset(torch.utils.data.Dataset):\n \"\"\" Real-world burst super-resolution dataset. \"\"\"\n def __init__(self, root, burst_size=8, crop_sz=80, center_crop=False, random_flip=False, split='train'):\n \"\"\"\n args:\n root : path of the root directory\n burst_size : Burst size. Maximum allowed burst size is 14.\n crop_sz: Size of the extracted crop. Maximum allowed crop size is 80\n center_crop: Whether to extract a random crop, or a centered crop.\n random_flip: Whether to apply random horizontal and vertical flip\n split: Can be 'train' or 'val'\n \"\"\"\n assert burst_size <= 14, 'burst_sz must be less than or equal to 14'\n assert crop_sz <= 80, 'crop_sz must be less than or equal to 80'\n assert split in ['train', 'val']\n\n root = root + '/' + split\n super().__init__()\n\n self.burst_size = burst_size\n self.crop_sz = crop_sz\n self.split = split\n self.center_crop = center_crop\n self.random_flip = random_flip\n\n self.root = root\n\n self.substract_black_level = True\n self.white_balance = False\n\n self.burst_list = self._get_burst_list()\n\n def _get_burst_list(self):\n burst_list = sorted(os.listdir('{}'.format(self.root)))\n # print(burst_list)\n return burst_list\n\n def get_burst_info(self, burst_id):\n burst_info = {'burst_size': 14, 'burst_name': self.burst_list[burst_id]}\n return burst_info\n\n def _get_raw_image(self, burst_id, im_id):\n raw_image = SamsungRAWImage.load('{}/{}/samsung_{:02d}'.format(self.root, self.burst_list[burst_id], im_id))\n return raw_image\n\n def _get_gt_image(self, burst_id):\n canon_im = CanonImage.load('{}/{}/canon'.format(self.root, self.burst_list[burst_id]), split=self.split)\n return canon_im\n\n def get_burst(self, burst_id, im_ids, info=None):\n frames = [self._get_raw_image(burst_id, i) for i in im_ids]\n\n gt = self._get_gt_image(burst_id)\n if info is None:\n info = self.get_burst_info(burst_id)\n\n return frames, gt, info\n\n def _sample_images(self):\n burst_size = 14\n\n ids = random.sample(range(1, burst_size), k=self.burst_size - 1)\n ids = [0, ] + ids\n return ids\n\n def __len__(self):\n return len(self.burst_list)\n\n def __getitem__(self, index):\n # Sample the images in the burst, in case a burst_size < 14 is used.\n im_ids = self._sample_images()\n\n # Read the burst images along with HR ground truth\n frames, gt, meta_info = self.get_burst(index, im_ids)\n\n # Extract crop if needed\n if frames[0].shape()[-1] != self.crop_sz:\n if getattr(self, 'center_crop', False):\n r1 = (frames[0].shape()[-2] - self.crop_sz) // 2\n c1 = (frames[0].shape()[-1] - self.crop_sz) // 2\n else:\n r1 = random.randint(0, frames[0].shape()[-2] - self.crop_sz)\n c1 = random.randint(0, frames[0].shape()[-1] - self.crop_sz)\n r2 = r1 + self.crop_sz\n c2 = c1 + self.crop_sz\n\n scale_factor = gt.shape()[-1] // frames[0].shape()[-1]\n frames = [im.get_crop(r1, r2, c1, c2) for im in frames]\n\n gt = gt.get_crop(scale_factor * r1, scale_factor * r2, scale_factor * c1, scale_factor * c2)\n\n # Load the RAW image data\n burst_image_data = [im.get_image_data(normalize=True, substract_black_level=self.substract_black_level,\n white_balance=self.white_balance) for im in frames]\n\n # Convert to tensor\n gt_image_data = gt.get_image_data(normalize=True, white_balance=self.white_balance,\n substract_black_level=self.substract_black_level)\n\n if self.random_flip:\n burst_image_data = [flatten_raw_image(im) for im in burst_image_data]\n\n pad = [0, 0, 0, 0]\n if random.random() > 0.5:\n burst_image_data = [im.flip([1, ])[:, 1:-1].contiguous() for im in burst_image_data]\n gt_image_data = gt_image_data.flip([2, ])[:, :, 2:-2].contiguous()\n pad[1] = 1\n\n if random.random() > 0.5:\n burst_image_data = [im.flip([0, ])[1:-1, :].contiguous() for im in burst_image_data]\n gt_image_data = gt_image_data.flip([1, ])[:, 2:-2, :].contiguous()\n pad[3] = 1\n\n burst_image_data = [pack_raw_image(im) for im in burst_image_data]\n burst_image_data = [F.pad(im.unsqueeze(0), pad, mode='replicate').squeeze(0) for im in burst_image_data]\n\n gt_image_data = F.pad(gt_image_data.unsqueeze(0), [4 * p for p in pad], mode='replicate').squeeze(0)\n\n burst_image_meta_info = frames[0].get_all_meta_data()\n\n burst_image_meta_info['black_level_subtracted'] = self.substract_black_level\n burst_image_meta_info['while_balance_applied'] = self.white_balance\n burst_image_meta_info['norm_factor'] = frames[0].norm_factor\n\n gt_image_meta_info = gt.get_all_meta_data()\n\n burst = torch.stack(burst_image_data, dim=0)\n\n burst_exposure = frames[0].get_exposure_time()\n canon_exposure = gt.get_exposure_time()\n\n burst_f_number = frames[0].get_f_number()\n canon_f_number = gt.get_f_number()\n\n burst_iso = frames[0].get_iso()\n canon_iso = gt.get_iso()\n\n # Normalize the GT image to account for differences in exposure, ISO etc\n light_factor_burst = burst_exposure * burst_iso / (burst_f_number ** 2)\n light_factor_canon = canon_exposure * canon_iso / (canon_f_number ** 2)\n\n exp_scale_factor = (light_factor_burst / light_factor_canon)\n gt_image_data = gt_image_data * exp_scale_factor\n\n gt_image_meta_info['black_level_subtracted'] = self.substract_black_level\n gt_image_meta_info['while_balance_applied'] = self.white_balance\n gt_image_meta_info['norm_factor'] = gt.norm_factor / exp_scale_factor\n\n burst_image_meta_info['exposure'] = burst_exposure\n burst_image_meta_info['f_number'] = burst_f_number\n burst_image_meta_info['iso'] = burst_iso\n\n gt_image_meta_info['exposure'] = canon_exposure\n gt_image_meta_info['f_number'] = canon_f_number\n gt_image_meta_info['iso'] = canon_iso\n\n burst = burst.float()\n frame_gt = gt_image_data.float()\n\n meta_info_burst = burst_image_meta_info\n meta_info_gt = gt_image_meta_info\n\n del meta_info_gt['crop_info']\n\n for k, v in meta_info_gt.items():\n if isinstance(v, (list, tuple)):\n meta_info_gt[k] = torch.tensor(v)\n\n for k, v in meta_info_burst.items():\n if isinstance(v, (list, tuple)):\n meta_info_burst[k] = torch.tensor(v)\n\n meta_info_burst['burst_name'] = meta_info['burst_name']\n\n return burst, frame_gt, meta_info_burst, meta_info_gt\n\n\ndef pack_raw_image(im_raw):\n if isinstance(im_raw, np.ndarray):\n im_out = np.zeros_like(im_raw, shape=(4, im_raw.shape[0] // 2, im_raw.shape[1] // 2))\n elif isinstance(im_raw, torch.Tensor):\n im_out = torch.zeros((4, im_raw.shape[0] // 2, im_raw.shape[1] // 2), dtype=im_raw.dtype).to(im_raw.device)\n else:\n raise Exception\n\n im_out[0, :, :] = im_raw[0::2, 0::2]\n im_out[1, :, :] = im_raw[0::2, 1::2]\n im_out[2, :, :] = im_raw[1::2, 0::2]\n im_out[3, :, :] = im_raw[1::2, 1::2]\n return im_out\n\n\ndef flatten_raw_image(im_raw_4ch):\n if isinstance(im_raw_4ch, np.ndarray):\n im_out = np.zeros_like(im_raw_4ch, shape=(im_raw_4ch.shape[1] * 2, im_raw_4ch.shape[2] * 2))\n elif isinstance(im_raw_4ch, torch.Tensor):\n im_out = torch.zeros((im_raw_4ch.shape[1] * 2, im_raw_4ch.shape[2] * 2), dtype=im_raw_4ch.dtype).to(im_raw_4ch.device)\n else:\n raise Exception\n\n im_out[0::2, 0::2] = im_raw_4ch[0, :, :]\n im_out[0::2, 1::2] = im_raw_4ch[1, :, :]\n im_out[1::2, 0::2] = im_raw_4ch[2, :, :]\n im_out[1::2, 1::2] = im_raw_4ch[3, :, :]\n\n return im_out\n\ndef pack_raw_image_batch(im_raw):\n im_out = torch.zeros((im_raw.shape[0], im_raw.shape[1], 4, im_raw.shape[3] // 2, im_raw.shape[4] // 2), dtype=im_raw.dtype).to(im_raw.device)\n im_out[:, :, 0, :, :] = im_raw[:, :, 0, 0::2, 0::2]\n im_out[:, :, 1, :, :] = im_raw[:, :, 0, 0::2, 1::2]\n im_out[:, :, 2, :, :] = im_raw[:, :, 0, 1::2, 0::2]\n im_out[:, :, 3, :, :] = im_raw[:, :, 0, 1::2, 1::2]\n return im_out\n\n\ndef flatten_raw_image_batch(im_raw_4ch):\n im_out = torch.zeros((im_raw_4ch.shape[0], im_raw_4ch.shape[1], 1, im_raw_4ch.shape[3] * 2, im_raw_4ch.shape[4] * 2), dtype=im_raw_4ch.dtype).to(im_raw_4ch.device)\n im_out[:, :, 0, 0::2, 0::2] = im_raw_4ch[:, :, 0, :, :]\n im_out[:, :, 0, 0::2, 1::2] = im_raw_4ch[:, :, 1, :, :]\n im_out[:, :, 0, 1::2, 0::2] = im_raw_4ch[:, :, 2, :, :]\n im_out[:, :, 0, 1::2, 1::2] = im_raw_4ch[:, :, 3, :, :]\n\n return im_out\n"
},
{
"path": "code/real/bsrt/datasets/burstsr_test_dataset.py",
"content": "import os\nimport torch\nimport torch.nn.functional as F\nimport random\nfrom .burstsr_dataset import SamsungRAWImage, flatten_raw_image, pack_raw_image\n\n\nclass BurstSRDataset(torch.utils.data.Dataset):\n \"\"\" Real-world burst super-resolution dataset. \"\"\"\n def __init__(self, root, burst_size=8, crop_sz=80, center_crop=False, random_flip=False, split='test'):\n \"\"\"\n args:\n root : path of the root directory\n burst_size : Burst size. Maximum allowed burst size is 14.\n crop_sz: Size of the extracted crop. Maximum allowed crop size is 80\n center_crop: Whether to extract a random crop, or a centered crop.\n random_flip: Whether to apply random horizontal and vertical flip\n split: Can be 'train' or 'val'\n \"\"\"\n assert burst_size <= 14, 'burst_sz must be less than or equal to 14'\n assert crop_sz <= 80, 'crop_sz must be less than or equal to 80'\n assert split in ['test']\n\n root = root + '/' + split\n super().__init__()\n\n self.burst_size = burst_size\n self.crop_sz = crop_sz\n self.split = split\n self.center_crop = center_crop\n self.random_flip = random_flip\n\n self.root = root\n\n self.substract_black_level = True\n self.white_balance = False\n\n self.burst_list = self._get_burst_list()\n\n def _get_burst_list(self):\n burst_list = sorted(os.listdir('{}'.format(self.root)))\n\n return burst_list\n\n def get_burst_info(self, burst_id):\n burst_info = {'burst_size': 14, 'burst_name': self.burst_list[burst_id]}\n return burst_info\n\n def _get_raw_image(self, burst_id, im_id):\n raw_image = SamsungRAWImage.load('{}/{}/samsung_{:02d}'.format(self.root, self.burst_list[burst_id], im_id))\n return raw_image\n\n def get_burst(self, burst_id, im_ids, info=None):\n frames = [self._get_raw_image(burst_id, i) for i in im_ids]\n\n if info is None:\n info = self.get_burst_info(burst_id)\n\n return frames, info\n\n def _sample_images(self):\n burst_size = 14\n\n ids = random.sample(range(1, burst_size), k=self.burst_size - 1)\n ids = [0, ] + ids\n return ids\n\n def __len__(self):\n return len(self.burst_list)\n\n def __getitem__(self, index):\n # Sample the images in the burst, in case a burst_size < 14 is used.\n im_ids = self._sample_images()\n\n # Read the burst images along with HR ground truth\n frames, meta_info = self.get_burst(index, im_ids)\n\n # Extract crop if needed\n if frames[0].shape()[-1] != self.crop_sz:\n if getattr(self, 'center_crop', False):\n r1 = (frames[0].shape()[-2] - self.crop_sz) // 2\n c1 = (frames[0].shape()[-1] - self.crop_sz) // 2\n else:\n r1 = random.randint(0, frames[0].shape()[-2] - self.crop_sz)\n c1 = random.randint(0, frames[0].shape()[-1] - self.crop_sz)\n r2 = r1 + self.crop_sz\n c2 = c1 + self.crop_sz\n\n frames = [im.get_crop(r1, r2, c1, c2) for im in frames]\n\n # Load the RAW image data\n burst_image_data = [im.get_image_data(normalize=True, substract_black_level=self.substract_black_level,\n white_balance=self.white_balance) for im in frames]\n\n if self.random_flip:\n burst_image_data = [flatten_raw_image(im) for im in burst_image_data]\n\n pad = [0, 0, 0, 0]\n if random.random() > 0.5:\n burst_image_data = [im.flip([1, ])[:, 1:-1].contiguous() for im in burst_image_data]\n pad[1] = 1\n\n if random.random() > 0.5:\n burst_image_data = [im.flip([0, ])[1:-1, :].contiguous() for im in burst_image_data]\n pad[3] = 1\n\n burst_image_data = [pack_raw_image(im) for im in burst_image_data]\n burst_image_data = [F.pad(im.unsqueeze(0), pad, mode='replicate').squeeze(0) for im in burst_image_data]\n\n burst_image_meta_info = frames[0].get_all_meta_data()\n\n burst_image_meta_info['black_level_subtracted'] = self.substract_black_level\n burst_image_meta_info['while_balance_applied'] = self.white_balance\n burst_image_meta_info['norm_factor'] = frames[0].norm_factor\n\n burst = torch.stack(burst_image_data, dim=0)\n\n burst_exposure = frames[0].get_exposure_time()\n\n burst_f_number = frames[0].get_f_number()\n\n burst_iso = frames[0].get_iso()\n\n burst_image_meta_info['exposure'] = burst_exposure\n burst_image_meta_info['f_number'] = burst_f_number\n burst_image_meta_info['iso'] = burst_iso\n\n burst = burst.float()\n\n meta_info_burst = burst_image_meta_info\n\n for k, v in meta_info_burst.items():\n if isinstance(v, (list, tuple)):\n meta_info_burst[k] = torch.tensor(v)\n\n return burst, meta_info_burst"
},
{
"path": "code/real/bsrt/datasets/data_sampler.py",
"content": "\"\"\"\nModified from torch.utils.data.distributed.DistributedSampler\nSupport enlarging the dataset for *iter-oriented* training, for saving time when restart the\ndataloader after each epoch\n\"\"\"\nimport math\n\nimport torch\nimport torch.distributed as dist\nfrom torch.utils.data.sampler import Sampler\n\n\nclass DistIterSampler(Sampler):\n \"\"\"Sampler that restricts data loading to a subset of the dataset.\n\n It is especially useful in conjunction with\n :class:`torch.nn.parallel.DistributedDataParallel`. In such case, each\n process can pass a DistributedSampler instance as a DataLoader sampler,\n and load a subset of the original dataset that is exclusive to it.\n\n .. note::\n Dataset is assumed to be of constant size.\n\n Arguments:\n dataset: Dataset used for sampling.\n num_replicas (optional): Number of processes participating in\n distributed training.\n rank (optional): Rank of the current process within num_replicas.\n \"\"\"\n\n def __init__(self, dataset, num_replicas=None, rank=None, ratio=100):\n if num_replicas is None:\n if not dist.is_available():\n raise RuntimeError(\"Requires distributed package to be available\")\n num_replicas = dist.get_world_size()\n if rank is None:\n if not dist.is_available():\n raise RuntimeError(\"Requires distributed package to be available\")\n rank = dist.get_rank()\n self.dataset = dataset\n self.num_replicas = num_replicas\n self.rank = rank\n self.epoch = 0\n self.num_samples = int(math.ceil(len(self.dataset) * ratio / self.num_replicas))\n self.total_size = self.num_samples * self.num_replicas\n\n def __iter__(self):\n # deterministically shuffle based on epoch\n g = torch.Generator()\n g.manual_seed(self.epoch)\n indices = torch.randperm(\n self.total_size, generator=g\n ).tolist() # Returns a random permutation of integers from 0 to n - 1\n\n dsize = len(self.dataset)\n indices = [v % dsize for v in indices]\n\n # subsample\n indices = indices[self.rank : self.total_size : self.num_replicas]\n assert len(indices) == self.num_samples\n\n return iter(indices)\n\n def __len__(self):\n return self.num_samples\n\n def set_epoch(self, epoch):\n self.epoch = epoch\n"
},
{
"path": "code/real/bsrt/datasets/realworld_burst_test_set.py",
"content": "import torch\nimport cv2\nimport numpy as np\nimport pickle as pkl\n\n\nclass RealWorldBurstTest(torch.utils.data.Dataset):\n \"\"\"\n \"\"\"\n def __init__(self, root):\n self.root = root\n self.burst_list = list(range(20))\n self.burst_size = 14\n\n def __len__(self):\n return len(self.burst_list)\n\n def _read_burst_image(self, index, image_id):\n im = cv2.imread('{}/{:04d}/im_raw_{:02d}.png'.format(self.root, index, image_id), cv2.IMREAD_UNCHANGED)\n im_t = torch.from_numpy(im.astype(np.float32)).permute(2, 0, 1).float() / (2**14)\n return im_t\n\n def __getitem__(self, index):\n \"\"\"\n args:\n index: Index of the burst\n\n returns:\n burst: LR RAW burst, a torch tensor of shape\n The 4 channels correspond to 'R', 'G', 'G', and 'B' values in the RGGB bayer mosaick.\n meta_info: Meta information about the burst\n \"\"\"\n burst_name = '{:04d}'.format(index)\n burst = [self._read_burst_image(index, i) for i in range(self.burst_size)]\n burst = torch.stack(burst, 0)\n\n meta_info = {}\n meta_info['burst_name'] = burst_name\n\n return burst, meta_info\n"
},
{
"path": "code/real/bsrt/datasets/synthetic_burst_test_set.py",
"content": "import torch\nimport cv2\nimport numpy as np\nimport pickle as pkl\n\n\nclass SyntheticBurstTest(torch.utils.data.Dataset):\n \"\"\" Synthetic burst test set. The test burst have been generated using the same synthetic pipeline as\n employed in SyntheticBurst dataset.\n \"\"\"\n def __init__(self, root):\n self.root = root\n self.burst_list = list(range(92))\n self.burst_size = 14\n\n def __len__(self):\n return len(self.burst_list)\n\n def _read_burst_image(self, index, image_id):\n im = cv2.imread('{}/{:04d}/im_raw_{:02d}.png'.format(self.root, index, image_id), cv2.IMREAD_UNCHANGED)\n im_t = torch.from_numpy(im.astype(np.float32)).permute(2, 0, 1).float() / (2**14)\n return im_t\n\n def __getitem__(self, index):\n \"\"\" Generates a synthetic burst\n args:\n index: Index of the burst\n\n returns:\n burst: LR RAW burst, a torch tensor of shape\n The 4 channels correspond to 'R', 'G', 'G', and 'B' values in the RGGB bayer mosaick.\n meta_info: Meta information about the burst\n \"\"\"\n burst_name = '{:04d}'.format(index)\n burst = [self._read_burst_image(index, i) for i in range(self.burst_size)]\n burst = torch.stack(burst, 0)\n\n meta_info = {}\n meta_info['burst_name'] = burst_name\n\n return burst, meta_info\n"
},
{
"path": "code/real/bsrt/datasets/synthetic_burst_train_set.py",
"content": "import torch\nimport numpy as np\nfrom PIL import Image\nfrom data_processing.synthetic_burst_generation import rgb2rawburst, random_crop #syn_burst_utils\nimport torchvision.transforms as tfm\n\n\nclass SyntheticBurst(torch.utils.data.Dataset):\n \"\"\" Synthetic burst dataset for joint denoising, demosaicking, and super-resolution. RAW Burst sequences are\n synthetically generated on the fly as follows. First, a single image is loaded from the base_dataset. The sampled\n image is converted to linear sensor space using the inverse camera pipeline employed in [1]. A burst\n sequence is then generated by adding random translations and rotations to the converted image. The generated burst\n is then converted is then mosaicked, and corrupted by random noise to obtain the RAW burst.\n\n [1] Unprocessing Images for Learned Raw Denoising, Brooks, Tim and Mildenhall, Ben and Xue, Tianfan and Chen,\n Jiawen and Sharlet, Dillon and Barron, Jonathan T, CVPR 2019\n \"\"\"\n def __init__(self, base_dataset, burst_size=8, crop_sz=384, transform=tfm.ToTensor()):\n self.base_dataset = base_dataset\n\n self.burst_size = burst_size\n self.crop_sz = crop_sz\n self.transform = transform\n\n self.downsample_factor = 4\n self.burst_transformation_params = {'max_translation': 24.0,\n 'max_rotation': 1.0,\n 'max_shear': 0.0,\n 'max_scale': 0.0,\n 'border_crop': 24}\n\n self.image_processing_params = {'random_ccm': True, 'random_gains': True, 'smoothstep': True,\n 'gamma': True,\n 'add_noise': True}\n self.interpolation_type = 'bilinear'\n\n def __len__(self):\n return len(self.base_dataset)\n\n def __getitem__(self, index):\n \"\"\" Generates a synthetic burst\n args:\n index: Index of the image in the base_dataset used to generate the burst\n\n returns:\n burst: Generated LR RAW burst, a torch tensor of shape\n [burst_size, 4, self.crop_sz / (2*self.downsample_factor), self.crop_sz / (2*self.downsample_factor)]\n The 4 channels correspond to 'R', 'G', 'G', and 'B' values in the RGGB bayer mosaick.\n The extra factor 2 in the denominator (2*self.downsample_factor) corresponds to the mosaicking\n operation.\n\n frame_gt: The HR RGB ground truth in the linear sensor space, a torch tensor of shape\n [3, self.crop_sz, self.crop_sz]\n\n flow_vectors: The ground truth flow vectors between a burst image and the base image (i.e. the first image in the burst).\n The flow_vectors can be used to warp the burst images to the base frame, using the 'warp'\n function in utils.warp package.\n flow_vectors is torch tensor of shape\n [burst_size, 2, self.crop_sz / self.downsample_factor, self.crop_sz / self.downsample_factor].\n Note that the flow_vectors are in the LR RGB space, before mosaicking. Hence it has twice\n the number of rows and columns, compared to the output burst.\n\n NOTE: The flow_vectors are only available during training for the purpose of using any\n auxiliary losses if needed. The flow_vectors will NOT be provided for the bursts in the\n test set\n\n meta_info: A dictionary containing the parameters used to generate the synthetic burst.\n \"\"\"\n frame = self.base_dataset[index]\n\n # Augmentation, e.g. convert to tensor\n if self.transform is not None:\n # frame = Image.fromarray(frame)\n frame = self.transform(frame)\n\n # Extract a random crop from the image\n crop_sz = self.crop_sz + 2 * self.burst_transformation_params.get('border_crop', 0)\n frame_crop = random_crop(frame, crop_sz)\n\n # Generate RAW burst\n burst, frame_gt, burst_rgb, flow_vectors, meta_info = rgb2rawburst(frame_crop,\n self.burst_size,\n self.downsample_factor,\n burst_transformation_params=self.burst_transformation_params,\n image_processing_params=self.image_processing_params,\n interpolation_type=self.interpolation_type\n )\n\n if self.burst_transformation_params.get('border_crop') is not None:\n border_crop = self.burst_transformation_params.get('border_crop')\n frame_gt = frame_gt[:, border_crop:-border_crop, border_crop:-border_crop]\n\n return burst, frame_gt, flow_vectors, meta_info\n"
},
{
"path": "code/real/bsrt/datasets/synthetic_burst_val_set.py",
"content": "import os\nimport torch\nimport cv2\nimport numpy as np\nimport pickle as pkl\n\n\nclass SyntheticBurstVal(torch.utils.data.Dataset):\n \"\"\" Synthetic burst validation set introduced in [1]. The validation burst have been generated using a\n synthetic data generation pipeline. The dataset can be downloaded from\n https://data.vision.ee.ethz.ch/bhatg/SyntheticBurstVal.zip\n\n [1] Deep Burst Super-Resolution. Goutam Bhat, Martin Danelljan, Luc Van Gool, and Radu Timofte. CVPR 2021\n \"\"\"\n def __init__(self, root=None, initialize=True):\n \"\"\"\n args:\n root - Path to root dataset directory\n initialize - boolean indicating whether to load the meta-data for the dataset\n \"\"\"\n self.root = os.path.join(root, 'val')\n self.burst_list = list(range(300))\n self.burst_size = 14\n\n def initialize(self):\n pass\n\n def __len__(self):\n return len(self.burst_list)\n\n def _read_burst_image(self, index, image_id):\n im = cv2.imread('{}/bursts/{:04d}/im_raw_{:02d}.png'.format(self.root, index, image_id), cv2.IMREAD_UNCHANGED)\n im_t = torch.from_numpy(im.astype(np.float32)).permute(2, 0, 1).float() / (2**14)\n\n return im_t\n\n def _read_gt_image(self, index):\n gt = cv2.imread('{}/gt/{:04d}/im_rgb.png'.format(self.root, index), cv2.IMREAD_UNCHANGED)\n gt_t = (torch.from_numpy(gt.astype(np.float32)) / 2 ** 14).permute(2, 0, 1).float()\n return gt_t\n\n def _read_meta_info(self, index):\n with open('{}/gt/{:04d}/meta_info.pkl'.format(self.root, index), \"rb\") as input_file:\n meta_info = pkl.load(input_file)\n\n return meta_info\n\n def __getitem__(self, index):\n \"\"\" Generates a synthetic burst\n args:\n index: Index of the burst\n\n returns:\n burst: LR RAW burst, a torch tensor of shape\n [14, 4, 48, 48]\n The 4 channels correspond to 'R', 'G', 'G', and 'B' values in the RGGB bayer mosaick.\n gt : Ground truth linear image\n meta_info: Meta info about the burst which can be used to convert gt to sRGB space\n \"\"\"\n burst_name = '{:04d}'.format(index)\n burst = [self._read_burst_image(index, i) for i in range(self.burst_size)]\n burst = torch.stack(burst, 0)\n\n gt = self._read_gt_image(index)\n meta_info = self._read_meta_info(index)\n meta_info['burst_name'] = burst_name\n return burst, gt, meta_info\n"
},
{
"path": "code/real/bsrt/datasets/zurich_raw2rgb_dataset.py",
"content": "import torch\nimport os\nimport numpy as np\nfrom cv2 import imread\n\n\nclass ZurichRAW2RGB(torch.utils.data.Dataset):\n \"\"\" Canon RGB images from the \"Zurich RAW to RGB mapping\" dataset. You can download the full\n dataset (22 GB) from http://people.ee.ethz.ch/~ihnatova/pynet.html#dataset. Alternatively, you can only download the\n Canon RGB images (5.5 GB) from https://data.vision.ee.ethz.ch/bhatg/zurich-raw-to-rgb.zip\n \"\"\"\n def __init__(self, root, split='train'):\n super().__init__()\n\n if split in ['train', 'test']:\n self.img_pth = os.path.join(root, split, 'canon')\n else:\n raise Exception('Unknown split {}'.format(split))\n\n self.image_list = self._get_image_list(split)\n\n def _get_image_list(self, split):\n if split == 'train':\n image_list = ['{:d}.jpg'.format(i) for i in range(46839)]\n elif split == 'test':\n # image_list = ['{:d}.jpg'.format(int(i)) for i in np.linspace(1, 1200, 200)]\n image_list = ['{:d}.jpg'.format(i) for i in range(1200)]\n else:\n raise Exception\n\n return image_list\n\n def _get_image(self, im_id):\n path = os.path.join(self.img_pth, self.image_list[im_id])\n img = imread(path)\n return img\n\n def get_image(self, im_id):\n frame = self._get_image(im_id)\n\n return frame\n\n def __len__(self):\n return len(self.image_list)\n\n def __getitem__(self, index):\n frame = self._get_image(index)\n\n return frame\n"
},
{
"path": "code/real/bsrt/demo.sh",
"content": "#!/usr/bin/env bash\n\n\npython main.py --n_GPUs 8 --print_every 20 --lr 0.00004 --decay 40-80 --save bsrt_tiny --model BSRT --fp16 --model_level S --swinfeature --batch_size 8 --burst_size 14 --patch_size 80 --pre_train ../../synthetic/train_log/bsrt/real_models/bsrt_tiny/bsrt_best_epoch.pth \n# python main.py --n_GPUs 8 --print_every 20 --lr 0.00004 --decay 40-80 --save bsrt_large --model BSRT --fp16 --model_level L --swinfeature --batch_size 8 --burst_size 14 --patch_size 48 --pre_train ../../synthetic/train_log/bsrt/real_models/bsrt_large/bsrt_best_epoch.pth \n\n\n# python test_real.py --n_GPUs 1 --model BSRT --model_level S --swinfeature --batch_size 1 --burst_size 14 --patch_size 80 --pre_train ../train_log/bsrt/real_models/bsrt_tiny/bsrtbest_epoch.pth --root /data/dataset/ntire21/burstsr/real\n# python test_real.py --n_GPUs 1 --model BSRT --model_level L --swinfeature --batch_size 1 --burst_size 14 --patch_size 80 --pre_train ../train_log/bsrt/real_models/bsrt_large/bsrt_realworld.pth --root /data/dataset/ntire21/burstsr/real"
},
{
"path": "code/real/bsrt/loss/Charbonnier.py",
"content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass CharbonnierLoss(nn.Module):\n \"\"\"L1 charbonnier loss.\"\"\"\n\n def __init__(self, epsilon=1e-3, reduce=True):\n super(CharbonnierLoss, self).__init__()\n self.eps = epsilon * epsilon\n self.reduce = reduce\n\n def forward(self, X, Y):\n diff = torch.add(X, -Y)\n error = torch.sqrt(diff * diff + self.eps)\n if self.reduce:\n loss = torch.mean(error)\n else:\n loss = error\n return loss"
},
{
"path": "code/real/bsrt/loss/__init__.py",
"content": "import os\nfrom importlib import import_module\n\nimport matplotlib\nmatplotlib.use('Agg')\nimport matplotlib.pyplot as plt\n\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass Loss(nn.modules.loss._Loss):\n def __init__(self, args, ckp):\n super(Loss, self).__init__()\n if args.local_rank == 0:\n print('Preparing loss function:')\n\n self.n_GPUs = args.n_GPUs\n self.loss = []\n self.loss_module = nn.ModuleList()\n for loss in args.loss.split('+'):\n weight, loss_type = loss.split('*')\n if loss_type == 'MSE':\n loss_function = nn.MSELoss()\n elif loss_type == 'L1':\n loss_function = nn.L1Loss()\n elif loss_type.find('VGG') >= 0:\n module = import_module('loss.vgg')\n loss_function = getattr(module, 'VGG')(\n loss_type[3:],\n rgb_range=args.rgb_range\n )\n elif loss_type.find('GAN') >= 0:\n module = import_module('loss.adversarial')\n loss_function = getattr(module, 'Adversarial')(\n args,\n loss_type\n )\n elif loss_type == 'FILTER':\n module = import_module('loss.filter')\n loss_function = getattr(module, 'Filter')(args)\n elif loss_type == 'SSIM':\n module = import_module('loss.mssim')\n loss_function = getattr(module, 'SSIM')(args)\n elif loss_type == 'MSSSIM':\n module = import_module('loss.mssim')\n loss_function = getattr(module, 'MSSSIM')(args)\n\n self.loss.append({\n 'type': loss_type,\n 'weight': float(weight),\n 'function': loss_function}\n )\n if loss_type.find('GAN') >= 0:\n self.loss.append({'type': 'DIS', 'weight': 1, 'function': None})\n\n if len(self.loss) > 1:\n self.loss.append({'type': 'Total', 'weight': 0, 'function': None})\n\n for l in self.loss:\n if l['function'] is not None:\n if args.local_rank == 0:\n print('{:.3f} * {}'.format(l['weight'], l['type']))\n self.loss_module.append(l['function'])\n\n self.log = torch.Tensor()\n\n device = torch.device('cpu' if args.cpu else 'cuda')\n self.loss_module.to(device)\n if args.precision == 'half': self.loss_module.half()\n if not args.cpu and args.n_GPUs > 1:\n self.loss_module = nn.DataParallel(\n self.loss_module, range(args.n_GPUs)\n )\n\n if args.load != '': self.load(ckp.dir, cpu=args.cpu)\n\n def forward(self, sr, hr):\n losses = []\n for i, l in enumerate(self.loss):\n if l['function'] is not None:\n loss = l['function'](sr, hr)\n effective_loss = l['weight'] * loss\n losses.append(effective_loss)\n self.log[-1, i] += effective_loss.item()\n elif l['type'] == 'DIS':\n self.log[-1, i] += self.loss[i - 1]['function'].loss\n\n loss_sum = sum(losses)\n if len(self.loss) > 1:\n self.log[-1, -1] += loss_sum.item()\n\n return loss_sum\n\n def step(self):\n for l in self.get_loss_module():\n if hasattr(l, 'scheduler'):\n l.scheduler.step()\n\n def start_log(self):\n self.log = torch.cat((self.log, torch.zeros(1, len(self.loss))))\n\n def end_log(self, n_batches):\n self.log[-1].div_(n_batches)\n\n def display_loss(self, batch):\n n_samples = batch + 1\n log = []\n for l, c in zip(self.loss, self.log[-1]):\n log.append('[{}: {:.4f}]'.format(l['type'], c / n_samples))\n\n return ''.join(log)\n\n def plot_loss(self, apath, epoch):\n axis = np.linspace(1, epoch, epoch)\n for i, l in enumerate(self.loss):\n label = '{} Loss'.format(l['type'])\n fig = plt.figure()\n plt.title(label)\n plt.plot(axis, self.log[:, i].numpy(), label=label)\n plt.legend()\n plt.xlabel('Epochs')\n plt.ylabel('Loss')\n plt.grid(True)\n plt.savefig(os.path.join(apath, 'loss_{}.pdf'.format(l['type'])))\n plt.close(fig)\n\n def get_loss_module(self):\n if self.n_GPUs == 1:\n return self.loss_module\n else:\n return self.loss_module.module\n\n def save(self, apath):\n torch.save(self.state_dict(), os.path.join(apath, 'loss.pt'))\n torch.save(self.log, os.path.join(apath, 'loss_log.pt'))\n\n def load(self, apath, cpu=False):\n if cpu:\n kwargs = {'map_location': lambda storage, loc: storage}\n else:\n kwargs = {}\n\n self.load_state_dict(torch.load(\n os.path.join(apath, 'loss.pt'),\n **kwargs\n ))\n self.log = torch.load(os.path.join(apath, 'loss_log.pt'))\n for l in self.get_loss_module():\n if hasattr(l, 'scheduler'):\n for _ in range(len(self.log)): l.scheduler.step()\n\n"
},
{
"path": "code/real/bsrt/loss/adversarial.py",
"content": "import utility\nfrom types import SimpleNamespace\n\nfrom model import common\nfrom loss import discriminator\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\n\nclass Adversarial(nn.Module):\n def __init__(self, args, gan_type):\n super(Adversarial, self).__init__()\n self.gan_type = gan_type\n self.gan_k = args.gan_k\n self.dis = discriminator.Discriminator(args)\n # if gan_type == 'WGAN_GP':\n if True:\n # see https://arxiv.org/pdf/1704.00028.pdf pp.4\n optim_dict = {\n 'optimizer': 'ADAM',\n 'betas': (0.5, 0.9),\n 'epsilon': 1e-8,\n 'lr': 1e-5,\n 'weight_decay': args.weight_decay,\n 'decay': args.decay,\n 'gamma': args.gamma\n }\n optim_args = SimpleNamespace(**optim_dict)\n else:\n optim_args = args\n\n self.optimizer = utility.make_optimizer(optim_args, self.dis)\n\n def forward(self, fake, real):\n # updating discriminator...\n self.loss = 0\n fake_detach = fake.detach() # do not backpropagate through G\n for _ in range(self.gan_k):\n self.optimizer.zero_grad()\n # d: B x 1 tensor\n d_fake = self.dis(fake_detach)\n d_real = self.dis(real)\n retain_graph = False\n if self.gan_type in ['GAN', 'SNGAN']:\n loss_d = self.bce(d_real, d_fake)\n elif self.gan_type.find('WGAN') >= 0:\n loss_d = (d_fake - d_real).mean()\n if self.gan_type.find('GP') >= 0:\n epsilon = torch.rand_like(fake).view(-1, 1, 1, 1)\n hat = fake_detach.mul(1 - epsilon) + real.mul(epsilon)\n hat.requires_grad = True\n d_hat = self.dis(hat)\n gradients = torch.autograd.grad(\n outputs=d_hat.sum(), inputs=hat,\n retain_graph=True, create_graph=True, only_inputs=True\n )[0]\n gradients = gradients.view(gradients.size(0), -1)\n gradient_norm = gradients.norm(2, dim=1)\n gradient_penalty = 10 * gradient_norm.sub(1).pow(2).mean()\n loss_d += gradient_penalty\n # from ESRGAN: Enhanced Super-Resolution Generative Adversarial Networks\n elif self.gan_type == 'RGAN':\n better_real = d_real - d_fake.mean(dim=0, keepdim=True)\n better_fake = d_fake - d_real.mean(dim=0, keepdim=True)\n loss_d = self.bce(better_real, better_fake)\n retain_graph = True\n\n # Discriminator update\n self.loss += loss_d.item()\n loss_d.backward(retain_graph=retain_graph)\n self.optimizer.step()\n\n if self.gan_type == 'WGAN':\n for p in self.dis.parameters():\n p.data.clamp_(-1, 1)\n\n self.loss /= self.gan_k\n\n # updating generator...\n d_fake_bp = self.dis(fake) # for backpropagation, use fake as it is\n if self.gan_type in ['GAN', 'SNGAN']:\n label_real = torch.ones_like(d_fake_bp)\n loss_g = F.binary_cross_entropy_with_logits(d_fake_bp, label_real)\n elif self.gan_type.find('WGAN') >= 0:\n loss_g = -d_fake_bp.mean()\n elif self.gan_type == 'RGAN':\n better_real = d_real.detach() - d_fake_bp.mean(dim=0, keepdim=True)\n better_fake = d_fake_bp - d_real.mean(dim=0, keepdim=True).detach()\n loss_g = self.bce(better_fake, better_real)\n\n # Generator loss\n return loss_g\n\n def state_dict(self, *args, **kwargs):\n state_discriminator = self.dis.state_dict(*args, **kwargs)\n state_optimizer = self.optimizer.state_dict()\n\n return dict(**state_discriminator, **state_optimizer)\n\n def bce(self, real, fake):\n label_real = torch.ones_like(real)\n label_fake = torch.zeros_like(fake)\n bce_real = F.binary_cross_entropy_with_logits(real, label_real)\n bce_fake = F.binary_cross_entropy_with_logits(fake, label_fake)\n bce_loss = bce_real + bce_fake\n return bce_loss\n\n# Some references\n# https://github.com/kuc2477/pytorch-wgan-gp/blob/master/model.py\n# OR\n# https://github.com/caogang/wgan-gp/blob/master/gan_cifar10.py\n"
},
{
"path": "code/real/bsrt/loss/discriminator.py",
"content": "from model import common\n\nimport torch.nn as nn\n\nclass Discriminator(nn.Module):\n '''\n output is not normalized\n '''\n def __init__(self, args, gan_type='GAN'):\n super(Discriminator, self).__init__()\n\n in_channels = args.n_colors\n out_channels = 32\n depth = 6\n\n def _block(_in_channels, _out_channels, stride=1):\n\n Conv = nn.Conv2d(\n _in_channels,\n _out_channels,\n 3,\n padding=1,\n stride=stride,\n bias=False\n )\n\n if gan_type == 'SNGAN':\n return nn.Sequential(\n spectral_norm(Conv),\n nn.BatchNorm2d(_out_channels),\n nn.LeakyReLU(negative_slope=0.2, inplace=True)\n )\n else:\n return nn.Sequential(\n Conv,\n nn.BatchNorm2d(_out_channels),\n nn.LeakyReLU(negative_slope=0.2, inplace=True)\n )\n\n m_features = [_block(in_channels, out_channels)]\n for i in range(depth):\n in_channels = out_channels\n # if i % 2 == 1:\n # stride = 1\n # out_channels *= 2\n # else:\n out_channels *= 2\n stride = 2\n m_features.append(_block(in_channels, out_channels, stride=stride))\n\n patch_size = args.patch_size // 2**(depth-1)\n\n # print(out_channels, patch_size)\n\n m_classifier = [\n nn.Flatten(),\n nn.Linear(out_channels*patch_size**2, 512),\n nn.LeakyReLU(0.2, True),\n nn.Linear(512, 1)\n ]\n\n self.features = nn.Sequential(*m_features)\n self.classifier = nn.Sequential(*m_classifier)\n\n def forward(self, x):\n features = self.features(x)\n # print(features.shape)\n output = self.classifier(features)\n\n return output\n\n"
},
{
"path": "code/real/bsrt/loss/filter.py",
"content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass Filter(nn.Module):\n def __init__(self, args):\n super().__init__()\n self.args = args\n\n kernel = torch.tensor([[1, 4, 1], [4, -20, 4], [1, 4, 1]])\n self.conv = nn.Conv2d(args.n_colors, args.n_colors, 3, 3)\n with torch.no_grad():\n self.conv.weight.copy_(kernel.float())\n self.loss = nn.L1Loss()\n\n def forward(self, x, y):\n preds_x = self.conv(x)\n preds_y = self.conv(y)\n\n return self.loss(preds_x, preds_y)\n"
},
{
"path": "code/real/bsrt/loss/hist_entropy.py",
"content": "import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\nclass HistEntropy(nn.Module):\n def __init__(self, args):\n super().__init__()\n self.args = args\n\n def forward(self, x):\n p = torch.softmax(x, dim=1)\n logp = torch.log_softmax(x, dim=1)\n\n entropy = (-p * logp).sum(dim=(2, 3)).mean()\n\n return entropy\n"
},
{
"path": "code/real/bsrt/loss/mssim.py",
"content": "import torch\nimport torch.nn.functional as F\nfrom math import exp\nimport numpy as np\n\n\ndef gaussian(window_size, sigma):\n gauss = torch.Tensor([exp(-(x - window_size//2)**2/float(2*sigma**2)) for x in range(window_size)])\n return gauss/gauss.sum()\n\n\ndef create_window(window_size, channel=1):\n _1D_window = gaussian(window_size, 1.5).unsqueeze(1)\n _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)\n window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()\n return window\n\n\ndef ssim(img1, img2, window_size=11, window=None, size_average=True, full=False, val_range=None):\n # Value range can be different from 255. Other common ranges are 1 (sigmoid) and 2 (tanh).\n if val_range is None:\n if torch.max(img1) > 128:\n max_val = 255\n else:\n max_val = 1\n\n if torch.min(img1) < -0.5:\n min_val = -1\n else:\n min_val = 0\n L = max_val - min_val\n else:\n L = val_range\n\n padd = 0\n (_, channel, height, width) = img1.size()\n if window is None:\n real_size = min(window_size, height, width)\n window = create_window(real_size, channel=channel).to(img1.device)\n\n mu1 = F.conv2d(img1, window, padding=padd, groups=channel)\n mu2 = F.conv2d(img2, window, padding=padd, groups=channel)\n\n mu1_sq = mu1.pow(2)\n mu2_sq = mu2.pow(2)\n mu1_mu2 = mu1 * mu2\n\n sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq\n sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq\n sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2\n\n C1 = (0.01 * L) ** 2\n C2 = (0.03 * L) ** 2\n\n v1 = 2.0 * sigma12 + C2\n v2 = sigma1_sq + sigma2_sq + C2\n cs = torch.mean(v1 / v2) # contrast sensitivity\n\n ssim_map = ((2 * mu1_mu2 + C1) * v1) / ((mu1_sq + mu2_sq + C1) * v2)\n\n if size_average:\n ret = ssim_map.mean()\n else:\n ret = ssim_map.mean(1).mean(1).mean(1)\n\n if full:\n return ret, cs\n return ret\n\n\ndef msssim(img1, img2, window_size=11, size_average=True, val_range=None, normalize=None):\n device = img1.device\n weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333]).to(device)\n levels = weights.size()[0]\n ssims = []\n mcs = []\n for _ in range(levels):\n sim, cs = ssim(img1, img2, window_size=window_size, size_average=size_average, full=True, val_range=val_range)\n\n # Relu normalize (not compliant with original definition)\n if normalize == \"relu\":\n ssims.append(torch.relu(sim))\n mcs.append(torch.relu(cs))\n else:\n ssims.append(sim)\n mcs.append(cs)\n\n img1 = F.avg_pool2d(img1, (2, 2))\n img2 = F.avg_pool2d(img2, (2, 2))\n\n ssims = torch.stack(ssims)\n mcs = torch.stack(mcs)\n\n # Simple normalize (not compliant with original definition)\n # TODO: remove support for normalize == True (kept for backward support)\n if normalize == \"simple\" or normalize == True:\n ssims = (ssims + 1) / 2\n mcs = (mcs + 1) / 2\n\n pow1 = mcs ** weights\n pow2 = ssims ** weights\n\n # From Matlab implementation https://ece.uwaterloo.ca/~z70wang/research/iwssim/\n output = torch.prod(pow1[:-1] * pow2[-1])\n return output\n\n\n# Classes to re-use window\nclass SSIM(torch.nn.Module):\n def __init__(self, window_size=11, size_average=True, val_range=None):\n super(SSIM, self).__init__()\n self.window_size = window_size\n self.size_average = size_average\n self.val_range = val_range\n\n # Assume 1 channel for SSIM\n self.channel = 1\n self.window = create_window(window_size)\n\n def forward(self, img1, img2):\n (_, channel, _, _) = img1.size()\n\n if channel == self.channel and self.window.dtype == img1.dtype:\n window = self.window\n else:\n window = create_window(self.window_size, channel).to(img1.device).type(img1.dtype)\n self.window = window\n self.channel = channel\n\n return ssim(img1, img2, window=window, window_size=self.window_size, size_average=self.size_average)\n\nclass MSSSIM(torch.nn.Module):\n def __init__(self, window_size=11, size_average=True, channel=3):\n super(MSSSIM, self).__init__()\n self.window_size = window_size\n self.size_average = size_average\n self.channel = channel\n\n def forward(self, img1, img2):\n # TODO: store window between calls if possible\n return msssim(img1, img2, window_size=self.window_size, size_average=self.size_average)"
},
{
"path": "code/real/bsrt/loss/vgg.py",
"content": "from model import common\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torchvision.models as models\n\nclass VGG(nn.Module):\n def __init__(self, conv_index, rgb_range=1):\n super(VGG, self).__init__()\n vgg_features = models.vgg19(pretrained=True).features\n modules = [m for m in vgg_features]\n if conv_index.find('22') >= 0:\n self.vgg = nn.Sequential(*modules[:8])\n elif conv_index.find('54') >= 0:\n self.vgg = nn.Sequential(*modules[:35])\n\n vgg_mean = (0.485, 0.456, 0.406)\n vgg_std = (0.229 * rgb_range, 0.224 * rgb_range, 0.225 * rgb_range)\n self.sub_mean = common.MeanShift(rgb_range, vgg_mean, vgg_std)\n for p in self.parameters():\n p.requires_grad = False\n\n def forward(self, sr, hr):\n def _forward(x):\n # x = self.sub_mean(x)\n x = self.vgg(x)\n return x\n\n sr = sr.repeat(1, 3, 1, 1)\n hr = hr.repeat(1, 3, 1, 1)\n\n vgg_sr = _forward(sr)\n with torch.no_grad():\n vgg_hr = _forward(hr.detach())\n\n loss = F.mse_loss(vgg_sr, vgg_hr)\n\n return loss\n"
},
{
"path": "code/real/bsrt/main.py",
"content": "import torch\nimport random\nimport numpy as np\nfrom torch.utils.data import DataLoader\nimport os\n\nimport utility\nimport model\nimport loss\nfrom option import args\nfrom trainer import Trainer\nfrom datasets.burstsr_dataset import BurstSRDataset, flatten_raw_image\nimport torch.multiprocessing as mp\nimport torch.backends.cudnn as cudnn\nimport torch.distributed as dist\nimport torch.utils.data.distributed\n\n\ndef init_seeds(seed=0, cuda_deterministic=True):\n random.seed(seed)\n np.random.seed(seed)\n torch.manual_seed(seed)\n # Speed-reproducibility tradeoff https://pytorch.org/docs/stable/notes/randomness.html\n if cuda_deterministic: # slower, more reproducible\n cudnn.deterministic = True\n cudnn.benchmark = False\n else: # faster, less reproducible\n cudnn.deterministic = False\n cudnn.benchmark = True\n\ncheckpoint = utility.checkpoint(args)\n\ndef main():\n mp.spawn(main_worker, nprocs=args.n_GPUs, args=(args.n_GPUs, args))\n\n\ndef main_worker(local_rank, nprocs, args):\n # print(local_rank)\n if checkpoint.ok:\n args.local_rank = local_rank\n init_seeds(local_rank+1)\n cudnn.benchmark = True\n utility.setup(local_rank, nprocs)\n torch.cuda.set_device(local_rank)\n\n batch_size = int(args.batch_size / nprocs)\n train_data = BurstSRDataset(root=args.root,\n burst_size=args.burst_size,\n crop_sz=args.patch_size, random_flip=True,\n center_crop=True, split='train')\n valid_data = BurstSRDataset(root=args.root,\n burst_size=14,\n crop_sz=80, split='val')\n\n if local_rank <= 0:\n print(f\"train data: {len(train_data)}, test data: {len(valid_data)}\")\n\n if nprocs > 1:\n train_sampler = torch.utils.data.distributed.DistributedSampler(train_data)\n valid_sampler = torch.utils.data.distributed.DistributedSampler(valid_data, shuffle=False)\n train_loader = DataLoader(dataset=train_data, batch_size=batch_size, num_workers=args.batch_size,\n pin_memory=True, drop_last=True, sampler=train_sampler) # args.cpus\n valid_loader = DataLoader(dataset=valid_data, batch_size=batch_size, num_workers=args.batch_size,\n pin_memory=True, drop_last=True, sampler=valid_sampler) # args.cpus\n else:\n train_sampler = None\n train_loader = DataLoader(dataset=train_data, batch_size=args.batch_size, num_workers=8,\n shuffle=True, pin_memory=True, drop_last=True) # args.cpus\n valid_loader = DataLoader(dataset=valid_data, batch_size=args.batch_size, num_workers=4, shuffle=False,\n pin_memory=True, drop_last=True) # args.cpus\n\n _model = model.Model(args, checkpoint)\n\n _loss = loss.Loss(args, checkpoint) if not args.test_only else None\n t = Trainer(args, train_loader, train_sampler, valid_loader, _model, _loss, checkpoint)\n while not t.terminate():\n t.train()\n\n del _model\n del _loss\n del train_loader\n del valid_loader\n\n # checkpoint.done()\n\nif __name__ == '__main__':\n # if not args.cpu: torch.cuda.set_device(0)\n main()\n"
},
{
"path": "code/real/bsrt/model/DCNv2/LICENSE",
"content": "BSD 3-Clause License\n\nCopyright (c) 2019, Charles Shang\nAll rights reserved.\n\nRedistribution and use in source and binary forms, with or without\nmodification, are permitted provided that the following conditions are met:\n\n1. Redistributions of source code must retain the above copyright notice, this\n list of conditions and the following disclaimer.\n\n2. Redistributions in binary form must reproduce the above copyright notice,\n this list of conditions and the following disclaimer in the documentation\n and/or other materials provided with the distribution.\n\n3. Neither the name of the copyright holder nor the names of its\n contributors may be used to endorse or promote products derived from\n this software without specific prior written permission.\n\nTHIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\nAND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\nIMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\nDISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE\nFOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL\nDAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR\nSERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER\nCAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,\nOR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE\nOF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE."
},
{
"path": "code/real/bsrt/model/DCNv2/README.md",
"content": "## Deformable Convolutional Networks V2 with Pytorch 1.0\n\n### Build\n```bash\n ./make.sh # build\n python test.py # run examples and gradient check \n```\n\n### An Example\n- deformable conv\n```python\n from dcn_v2 import DCN\n input = torch.randn(2, 64, 128, 128).cuda()\n # wrap all things (offset and mask) in DCN\n dcn = DCN(64, 64, kernel_size=(3,3), stride=1, padding=1, deformable_groups=2).cuda()\n output = dcn(input)\n print(output.shape)\n```\n- deformable roi pooling\n```python\n from dcn_v2 import DCNPooling\n input = torch.randn(2, 32, 64, 64).cuda()\n batch_inds = torch.randint(2, (20, 1)).cuda().float()\n x = torch.randint(256, (20, 1)).cuda().float()\n y = torch.randint(256, (20, 1)).cuda().float()\n w = torch.randint(64, (20, 1)).cuda().float()\n h = torch.randint(64, (20, 1)).cuda().float()\n rois = torch.cat((batch_inds, x, y, x + w, y + h), dim=1)\n\n # mdformable pooling (V2)\n # wrap all things (offset and mask) in DCNPooling\n dpooling = DCNPooling(spatial_scale=1.0 / 4,\n pooled_size=7,\n output_dim=32,\n no_trans=False,\n group_size=1,\n trans_std=0.1).cuda()\n\n dout = dpooling(input, rois)\n```\n### Note\nNow the master branch is for pytorch 1.0 (new ATen API), you can switch back to pytorch 0.4 with,\n```bash\ngit checkout pytorch_0.4\n```\n\n### Known Issues:\n\n- [x] Gradient check w.r.t offset (solved)\n- [ ] Backward is not reentrant (minor)\n\nThis is an adaption of the official [Deformable-ConvNets](https://github.com/msracver/Deformable-ConvNets/tree/master/DCNv2_op).\n\nI have ran the gradient check for many times with DOUBLE type. Every tensor **except offset** passes.\nHowever, when I set the offset to 0.5, it passes. I'm still wondering what cause this problem. Is it because some\nnon-differential points? \n\nUpdate: all gradient check passes with double precision. \n\nAnother issue is that it raises `RuntimeError: Backward is not reentrant`. However, the error is very small (`<1e-7` for \nfloat `<1e-15` for double), \nso it may not be a serious problem (?)\n\nPlease post an issue or PR if you have any comments.\n "
},
{
"path": "code/real/bsrt/model/DCNv2/__init__.py",
"content": ""
},
{
"path": "code/real/bsrt/model/DCNv2/dcn_v2.py",
"content": "#!/usr/bin/env python\nfrom __future__ import absolute_import, division, print_function\n\nimport math\n\nimport torch\nfrom torch import nn\nfrom torch.autograd import Function\nfrom torch.autograd.function import once_differentiable\nfrom torch.nn.modules.utils import _pair\nfrom torch.cuda.amp import custom_fwd, custom_bwd\n# from apex import amp\n\nimport _ext as _backend\n\n\nclass _DCNv2(Function):\n @staticmethod\n @custom_fwd(cast_inputs=torch.float32)\n # @amp.float_function\n def forward(\n ctx, input, offset, mask, weight, bias, stride, padding, dilation, deformable_groups\n ):\n ctx.stride = _pair(stride)\n ctx.padding = _pair(padding)\n ctx.dilation = _pair(dilation)\n ctx.kernel_size = _pair(weight.shape[2:4])\n ctx.deformable_groups = deformable_groups\n output = _backend.dcn_v2_forward(\n input,\n weight,\n bias,\n offset,\n mask,\n ctx.kernel_size[0],\n ctx.kernel_size[1],\n ctx.stride[0],\n ctx.stride[1],\n ctx.padding[0],\n ctx.padding[1],\n ctx.dilation[0],\n ctx.dilation[1],\n ctx.deformable_groups,\n )\n ctx.save_for_backward(input, offset, mask, weight, bias)\n return output\n\n @staticmethod\n @once_differentiable\n @custom_bwd\n # @amp.float_function\n def backward(ctx, grad_output):\n input, offset, mask, weight, bias = ctx.saved_tensors\n grad_input, grad_offset, grad_mask, grad_weight, grad_bias = _backend.dcn_v2_backward(\n input,\n weight,\n bias,\n offset,\n mask,\n grad_output,\n ctx.kernel_size[0],\n ctx.kernel_size[1],\n ctx.stride[0],\n ctx.stride[1],\n ctx.padding[0],\n ctx.padding[1],\n ctx.dilation[0],\n ctx.dilation[1],\n ctx.deformable_groups,\n )\n\n return grad_input, grad_offset, grad_mask, grad_weight, grad_bias, None, None, None, None\n\n @staticmethod\n def symbolic(\n g, input, offset, mask, weight, bias, stride, padding, dilation, deformable_groups\n ):\n from torch.nn.modules.utils import _pair\n\n stride = _pair(stride)\n padding = _pair(padding)\n dilation = _pair(dilation)\n # as of trt 7, the dcn operation will be translated again by modifying the onnx file\n # so the exporting code is kept to resemble the forward()\n return g.op(\n \"DCNv2_2\",\n input,\n offset,\n mask,\n weight,\n bias,\n stride_i=stride,\n padding_i=padding,\n dilation_i=dilation,\n deformable_groups_i=deformable_groups,\n )\n\n\ndcn_v2_conv = _DCNv2.apply\n\n\nclass DCNv2(nn.Module):\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n dilation=1,\n deformable_groups=1,\n ):\n super(DCNv2, self).__init__()\n self.in_channels = in_channels\n self.out_channels = out_channels\n self.kernel_size = _pair(kernel_size)\n self.stride = _pair(stride)\n self.padding = _pair(padding)\n self.dilation = _pair(dilation)\n self.deformable_groups = deformable_groups\n\n self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels, *self.kernel_size))\n self.bias = nn.Parameter(torch.Tensor(out_channels))\n self.reset_parameters()\n\n def reset_parameters(self):\n n = self.in_channels\n for k in self.kernel_size:\n n *= k\n stdv = 1.0 / math.sqrt(n)\n self.weight.data.uniform_(-stdv, stdv)\n self.bias.data.zero_()\n\n def forward(self, input, offset, mask):\n assert (\n 2 * self.deformable_groups * self.kernel_size[0] * self.kernel_size[1]\n == offset.shape[1]\n )\n assert self.deformable_groups * self.kernel_size[0] * self.kernel_size[1] == mask.shape[1]\n return dcn_v2_conv(\n input,\n offset,\n mask,\n self.weight,\n self.bias,\n self.stride,\n self.padding,\n self.dilation,\n self.deformable_groups,\n )\n\n\nclass DCN(DCNv2):\n def __init__(\n self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n dilation=1,\n deformable_groups=1,\n ):\n super(DCN, self).__init__(\n in_channels, out_channels, kernel_size, stride, padding, dilation, deformable_groups\n )\n\n channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]\n self.conv_offset_mask = nn.Conv2d(\n self.in_channels,\n channels_,\n kernel_size=self.kernel_size,\n stride=self.stride,\n padding=self.padding,\n bias=True,\n )\n self.init_offset()\n\n def init_offset(self):\n self.conv_offset_mask.weight.data.zero_()\n self.conv_offset_mask.bias.data.zero_()\n\n def forward(self, input):\n out = self.conv_offset_mask(input)\n o1, o2, mask = torch.chunk(out, 3, dim=1)\n offset = torch.cat((o1, o2), dim=1)\n mask = torch.sigmoid(mask)\n return dcn_v2_conv(\n input,\n offset,\n mask,\n self.weight,\n self.bias,\n self.stride,\n self.padding,\n self.dilation,\n self.deformable_groups,\n )\n\n\nclass DCN_sep(DCNv2):\n '''Use other features to generate offsets and masks'''\n\n def __init__(self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n dilation=1,\n deformable_groups=1):\n super(DCN_sep, self).__init__(in_channels, out_channels, kernel_size, stride, padding,\n dilation, deformable_groups)\n\n channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]\n self.conv_offset_mask = nn.Conv2d(\n self.in_channels,\n channels_,\n kernel_size=self.kernel_size,\n stride=self.stride,\n padding=self.padding,\n bias=True)\n self.init_offset()\n\n def init_offset(self):\n self.conv_offset_mask.weight.data.zero_()\n self.conv_offset_mask.bias.data.zero_()\n\n def forward(self, input, fea):\n '''input: input features for deformable conv\n fea: other features used for generating offsets and mask'''\n out = self.conv_offset_mask(fea)\n o1, o2, mask = torch.chunk(out, 3, dim=1)\n offset = torch.cat((o1, o2), dim=1)\n # offset = torch.clamp(offset, -100, 100)\n\n offset_mean = torch.mean(torch.abs(offset))\n if offset_mean > 250:\n print('Offset mean is {}, larger than 100.'.format(offset_mean))\n # return None\n # offset[offset>=150] = 1e-3\n # offset = offset.clamp(-50, 50)\n\n mask = torch.sigmoid(mask)\n return dcn_v2_conv(input, offset, mask, self.weight, self.bias, self.stride, self.padding,\n self.dilation, self.deformable_groups)\n\n\nclass FlowGuidedDCN(DCNv2):\n '''Use other features to generate offsets and masks'''\n\n def __init__(self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n dilation=1,\n deformable_groups=1):\n super(FlowGuidedDCN, self).__init__(in_channels, out_channels, kernel_size, stride, padding,\n dilation, deformable_groups)\n\n channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]\n self.conv_offset_mask = nn.Conv2d(\n in_channels, channels_, kernel_size, stride, padding, bias=True)\n\n self.init_offset()\n\n def init_offset(self):\n self.conv_offset_mask.weight.data.zero_()\n self.conv_offset_mask.bias.data.zero_()\n\n def forward(self, input, fea, flows):\n '''input: input features for deformable conv: N, C, H, W.\n fea: other features used for generating offsets and mask: N, C, H, W.\n flows: N, 2, H, W.\n '''\n out = self.conv_offset_mask(fea)\n o1, o2, mask = torch.chunk(out, 3, dim=1)\n\n offset = torch.tanh(torch.cat((o1, o2), dim=1)) * 10 # max_residue_magnitude\n offset = offset + flows.flip(1).repeat(1, offset.size(1)//2, 1, 1)\n\n offset_mean = torch.mean(torch.abs(offset))\n if offset_mean > 250:\n print('FlowGuidedDCN: Offset mean is {}, larger than 100.'.format(offset_mean))\n # offset = offset.clamp(-50, 50)\n # return None\n\n mask = torch.sigmoid(mask)\n return dcn_v2_conv(input, offset, mask, self.weight, self.bias, self.stride, self.padding,\n self.dilation, self.deformable_groups)\n\n\n\nclass InsideFlowGuidedDCN(DCNv2):\n '''Use other features to generate offsets and masks'''\n\n def __init__(self,\n in_channels,\n out_channels,\n kernel_size,\n stride,\n padding,\n dilation=1,\n deformable_groups=1):\n super(InsideFlowGuidedDCN, self).__init__(in_channels, out_channels, kernel_size, stride, padding,\n dilation, deformable_groups)\n\n channels_ = self.deformable_groups * 3 * self.kernel_size[0] * self.kernel_size[1]\n self.conv_offset_mask = nn.Sequential(\n nn.Conv2d(in_channels*2+2, out_channels, kernel_size, stride, padding, bias=True),\n nn.LeakyReLU(negative_slope=0.1, inplace=True),\n nn.Conv2d(out_channels, out_channels, kernel_size, stride, padding, bias=True),\n nn.LeakyReLU(negative_slope=0.1, inplace=True),\n nn.Conv2d(out_channels, channels_, kernel_size, stride, padding, bias=True)\n )\n\n self.reset_parameters()\n self.init_offset()\n\n def reset_parameters(self):\n n = self.in_channels\n for k in self.kernel_size:\n n *= k\n stdv = 1.0 / math.sqrt(n)\n self.weight.data.uniform_(-stdv, stdv)\n self.bias.data.zero_()\n\n\n def init_offset(self):\n self.conv_offset_mask[-1].weight.data.zero_()\n self.conv_offset_mask[-1].bias.data.zero_()\n\n def forward(self, input, warped, ref, flows):\n '''input: input features for deformable conv: N, C, H, W.\n fea: other features used for generating offsets and mask: N, C, H, W.\n flows: N, 2, H, W.\n '''\n out = self.conv_offset_mask(torch.cat([warped, ref, flows], dim=1))\n o1, o2, mask = torch.chunk(out, 3, dim=1)\n\n offset = torch.tanh(torch.cat((o1, o2), dim=1)) * 10 # max_residue_magnitude\n offset = offset + flows.flip(1).repeat(1, offset.size(1)//2, 1, 1)\n\n offset_mean = torch.mean(torch.abs(offset))\n if offset_mean > 250:\n print('InsideFlowGuidedDCN: Offset mean is {}, larger than 100.'.format(offset_mean))\n print('flow mean is {}'.format(torch.abs(flows).mean()))\n offset = offset.clamp(-50, 50)\n # return None\n\n mask = torch.sigmoid(mask)\n return dcn_v2_conv(input, offset, mask, self.weight, self.bias, self.stride, self.padding,\n self.dilation, self.deformable_groups)\n\n\n\nclass _DCNv2Pooling(Function):\n @staticmethod\n def forward(\n ctx,\n input,\n rois,\n offset,\n spatial_scale,\n pooled_size,\n output_dim,\n no_trans,\n group_size=1,\n part_size=None,\n sample_per_part=4,\n trans_std=0.0,\n ):\n ctx.spatial_scale = spatial_scale\n ctx.no_trans = int(no_trans)\n ctx.output_dim = output_dim\n ctx.group_size = group_size\n ctx.pooled_size = pooled_size\n ctx.part_size = pooled_size if part_size is None else part_size\n ctx.sample_per_part = sample_per_part\n ctx.trans_std = trans_std\n\n output, output_count = _backend.dcn_v2_psroi_pooling_forward(\n input,\n rois,\n offset,\n ctx.no_trans,\n ctx.spatial_scale,\n ctx.output_dim,\n ctx.group_size,\n ctx.pooled_size,\n ctx.part_size,\n ctx.sample_per_part,\n ctx.trans_std,\n )\n ctx.save_for_backward(input, rois, offset, output_count)\n return output\n\n @staticmethod\n @once_differentiable\n def backward(ctx, grad_output):\n input, rois, offset, output_count = ctx.saved_tensors\n grad_input, grad_offset = _backend.dcn_v2_psroi_pooling_backward(\n grad_output,\n input,\n rois,\n offset,\n output_count,\n ctx.no_trans,\n ctx.spatial_scale,\n ctx.output_dim,\n ctx.group_size,\n ctx.pooled_size,\n ctx.part_size,\n ctx.sample_per_part,\n ctx.trans_std,\n )\n\n return grad_input, None, grad_offset, None, None, None, None, None, None, None, None\n\n\ndcn_v2_pooling = _DCNv2Pooling.apply\n\n\nclass DCNv2Pooling(nn.Module):\n def __init__(\n self,\n spatial_scale,\n pooled_size,\n output_dim,\n no_trans,\n group_size=1,\n part_size=None,\n sample_per_part=4,\n trans_std=0.0,\n ):\n super(DCNv2Pooling, self).__init__()\n self.spatial_scale = spatial_scale\n self.pooled_size = pooled_size\n self.output_dim = output_dim\n self.no_trans = no_trans\n self.group_size = group_size\n self.part_size = pooled_size if part_size is None else part_size\n self.sample_per_part = sample_per_part\n self.trans_std = trans_std\n\n def forward(self, input, rois, offset):\n assert input.shape[1] == self.output_dim\n if self.no_trans:\n offset = input.new()\n return dcn_v2_pooling(\n input,\n rois,\n offset,\n self.spatial_scale,\n self.pooled_size,\n self.output_dim,\n self.no_trans,\n self.group_size,\n self.part_size,\n self.sample_per_part,\n self.trans_std,\n )\n\n\nclass DCNPooling(DCNv2Pooling):\n def __init__(\n self,\n spatial_scale,\n pooled_size,\n output_dim,\n no_trans,\n group_size=1,\n part_size=None,\n sample_per_part=4,\n trans_std=0.0,\n deform_fc_dim=1024,\n ):\n super(DCNPooling, self).__init__(\n spatial_scale,\n pooled_size,\n output_dim,\n no_trans,\n group_size,\n part_size,\n sample_per_part,\n trans_std,\n )\n\n self.deform_fc_dim = deform_fc_dim\n\n if not no_trans:\n self.offset_mask_fc = nn.Sequential(\n nn.Linear(\n self.pooled_size * self.pooled_size * self.output_dim, self.deform_fc_dim\n ),\n nn.ReLU(inplace=True),\n nn.Linear(self.deform_fc_dim, self.deform_fc_dim),\n nn.ReLU(inplace=True),\n nn.Linear(self.deform_fc_dim, self.pooled_size * self.pooled_size * 3),\n )\n self.offset_mask_fc[4].weight.data.zero_()\n self.offset_mask_fc[4].bias.data.zero_()\n\n def forward(self, input, rois):\n offset = input.new()\n\n if not self.no_trans:\n\n # do roi_align first\n n = rois.shape[0]\n roi = dcn_v2_pooling(\n input,\n rois,\n offset,\n self.spatial_scale,\n self.pooled_size,\n self.output_dim,\n True, # no trans\n self.group_size,\n self.part_size,\n self.sample_per_part,\n self.trans_std,\n )\n\n # build mask and offset\n offset_mask = self.offset_mask_fc(roi.view(n, -1))\n offset_mask = offset_mask.view(n, 3, self.pooled_size, self.pooled_size)\n o1, o2, mask = torch.chunk(offset_mask, 3, dim=1)\n offset = torch.cat((o1, o2), dim=1)\n mask = torch.sigmoid(mask)\n\n # do pooling with offset and mask\n return (\n dcn_v2_pooling(\n input,\n rois,\n offset,\n self.spatial_scale,\n self.pooled_size,\n self.output_dim,\n self.no_trans,\n self.group_size,\n self.part_size,\n self.sample_per_part,\n self.trans_std,\n )\n * mask\n )\n # only roi_align\n return dcn_v2_pooling(\n input,\n rois,\n offset,\n self.spatial_scale,\n self.pooled_size,\n self.output_dim,\n self.no_trans,\n self.group_size,\n self.part_size,\n self.sample_per_part,\n self.trans_std,\n )\n"
},
{
"path": "code/real/bsrt/model/DCNv2/files.txt",
"content": "/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/_ext.cpython-37m-x86_64-linux-gnu.so\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/_ext.py\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/EGG-INFO/PKG-INFO\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/EGG-INFO/SOURCES.txt\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/EGG-INFO/dependency_links.txt\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/EGG-INFO/native_libs.txt\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/EGG-INFO/not-zip-safe\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/EGG-INFO/top_level.txt\n/home/luoziwei/miniconda3/lib/python3.7/site-packages/DCNv2-0.1-py3.7-linux-x86_64.egg/__pycache__/_ext.cpython-37.pyc\n"
},
{
"path": "code/real/bsrt/model/DCNv2/make.sh",
"content": "#!/usr/bin/env bash\npython setup.py build develop\n"
},
{
"path": "code/real/bsrt/model/DCNv2/setup.py",
"content": "#!/usr/bin/env python\n\nimport glob\nimport os\n\nimport torch\nfrom setuptools import find_packages, setup\nfrom torch.utils.cpp_extension import CUDA_HOME, CppExtension, CUDAExtension\n\nrequirements = [\"torch\", \"torchvision\"]\n\n\ndef get_extensions():\n this_dir = os.path.dirname(os.path.abspath(__file__))\n extensions_dir = os.path.join(this_dir, \"src\")\n\n main_file = glob.glob(os.path.join(extensions_dir, \"*.cpp\"))\n source_cpu = glob.glob(os.path.join(extensions_dir, \"cpu\", \"*.cpp\"))\n source_cuda = glob.glob(os.path.join(extensions_dir, \"cuda\", \"*.cu\"))\n\n os.environ[\"CC\"] = \"g++\"\n sources = main_file + source_cpu\n extension = CppExtension\n extra_compile_args = {\"cxx\": []}\n define_macros = []\n\n if True:\n extension = CUDAExtension\n sources += source_cuda\n define_macros += [(\"WITH_CUDA\", None)]\n extra_compile_args[\"nvcc\"] = [\n \"-DCUDA_HAS_FP16=1\",\n \"-D__CUDA_NO_HALF_OPERATORS__\",\n \"-D__CUDA_NO_HALF_CONVERSIONS__\",\n \"-D__CUDA_NO_HALF2_OPERATORS__\",\n ]\n else:\n # raise NotImplementedError('Cuda is not available')\n pass\n\n sources = [os.path.join(extensions_dir, s) for s in sources]\n include_dirs = [extensions_dir]\n ext_modules = [\n extension(\n \"_ext\",\n sources,\n include_dirs=include_dirs,\n define_macros=define_macros,\n extra_compile_args=extra_compile_args,\n )\n ]\n return ext_modules\n\n\nsetup(\n name=\"DCNv2\",\n version=\"0.1\",\n author=\"charlesshang\",\n url=\"https://github.com/charlesshang/DCNv2\",\n description=\"deformable convolutional networks\",\n packages=find_packages(exclude=(\"configs\", \"tests\")),\n # install_requires=requirements,\n ext_modules=get_extensions(),\n cmdclass={\"build_ext\": torch.utils.cpp_extension.BuildExtension},\n)\n"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cpu/dcn_v2_cpu.cpp",
"content": "#include \n#include \"cpu/dcn_v2_im2col_cpu.h\"\n\n#include \n//#include \n\n#include \n//#include \n//#include \n\n//extern THCState *state;\n\n// author: Charles Shang\n// https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu\n// modified from the CUDA version for CPU use by Daniel K. Suhendro\n\nat::Tensor\ndcn_v2_cpu_forward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const int kernel_h,\n const int kernel_w,\n const int stride_h,\n const int stride_w,\n const int pad_h,\n const int pad_w,\n const int dilation_h,\n const int dilation_w,\n const int deformable_group)\n{\n // THCAssertSameGPU(THCudaTensor_checkGPU(state, 5, input, weight, bias, offset, mask));\n /*AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(weight.type().is_cuda(), \"weight must be a CUDA tensor\");\n AT_ASSERTM(bias.type().is_cuda(), \"bias must be a CUDA tensor\");\n AT_ASSERTM(offset.type().is_cuda(), \"offset must be a CUDA tensor\");\n AT_ASSERTM(mask.type().is_cuda(), \"mask must be a CUDA tensor\");*/\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n\n const int channels_out = weight.size(0);\n const int channels_kernel = weight.size(1);\n const int kernel_h_ = weight.size(2);\n const int kernel_w_ = weight.size(3);\n\n // printf(\"Kernels: %d %d %d %d\\n\", kernel_h_, kernel_w_, kernel_w, kernel_h);\n // printf(\"Channels: %d %d\\n\", channels, channels_kernel);\n // printf(\"Channels: %d %d\\n\", channels_out, channels_kernel);\n\n AT_ASSERTM(kernel_h_ == kernel_h && kernel_w_ == kernel_w,\n \"Input shape and kernel shape wont match: (%d x %d vs %d x %d).\", kernel_h_, kernel_w, kernel_h_, kernel_w_);\n\n AT_ASSERTM(channels == channels_kernel,\n \"Input shape and kernel channels wont match: (%d vs %d).\", channels, channels_kernel);\n\n const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;\n const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;\n\n auto ones = at::ones({height_out, width_out}, input.options());\n auto columns = at::empty({channels * kernel_h * kernel_w, 1 * height_out * width_out}, input.options());\n auto output = at::empty({batch, channels_out, height_out, width_out}, input.options());\n\n using scalar_t = float;\n for (int b = 0; b < batch; b++)\n {\n auto input_n = input.select(0, b);\n auto offset_n = offset.select(0, b);\n auto mask_n = mask.select(0, b);\n auto output_n = output.select(0, b);\n\n // Do Bias first:\n // M,N,K are dims of matrix A and B\n // (see http://docs.nvidia.com/cuda/cublas/#cublas-lt-t-gt-gemm)\n // (N x 1) (1 x M)\n long m_ = channels_out;\n long n_ = height_out * width_out;\n long k_ = 1;\n THFloatBlas_gemm('t', 'n', n_, m_, k_, 1.0f,\n ones.contiguous().data(), k_,\n bias.contiguous().data(), k_, 0.0f,\n output_n.data(), n_);\n\n modulated_deformable_im2col_cpu(input_n.data(),\n offset_n.data(),\n mask_n.data(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,\n deformable_group,\n columns.data());\n\n //(k * m) x (m * n)\n // Y = WC\n long m = channels_out;\n long n = height_out * width_out;\n long k = channels * kernel_h * kernel_w;\n THFloatBlas_gemm('n', 'n', n, m, k, 1.0f,\n columns.data(), n,\n weight.data(), k, 1.0f,\n output_n.data(), n);\n }\n return output;\n}\n\nstd::vector dcn_v2_cpu_backward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const at::Tensor &grad_output,\n int kernel_h, int kernel_w,\n int stride_h, int stride_w,\n int pad_h, int pad_w,\n int dilation_h, int dilation_w,\n int deformable_group)\n{\n\n THArgCheck(input.is_contiguous(), 1, \"input tensor has to be contiguous\");\n THArgCheck(weight.is_contiguous(), 2, \"weight tensor has to be contiguous\");\n\n /*AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(weight.type().is_cuda(), \"weight must be a CUDA tensor\");\n AT_ASSERTM(bias.type().is_cuda(), \"bias must be a CUDA tensor\");\n AT_ASSERTM(offset.type().is_cuda(), \"offset must be a CUDA tensor\");\n AT_ASSERTM(mask.type().is_cuda(), \"mask must be a CUDA tensor\");*/\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n\n const int channels_out = weight.size(0);\n const int channels_kernel = weight.size(1);\n const int kernel_h_ = weight.size(2);\n const int kernel_w_ = weight.size(3);\n\n AT_ASSERTM(kernel_h_ == kernel_h && kernel_w_ == kernel_w,\n \"Input shape and kernel shape wont match: (%d x %d vs %d x %d).\", kernel_h_, kernel_w, kernel_h_, kernel_w_);\n\n AT_ASSERTM(channels == channels_kernel,\n \"Input shape and kernel channels wont match: (%d vs %d).\", channels, channels_kernel);\n\n const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;\n const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;\n\n auto ones = at::ones({height_out, width_out}, input.options());\n auto columns = at::empty({channels * kernel_h * kernel_w, 1 * height_out * width_out}, input.options());\n auto output = at::empty({batch, channels_out, height_out, width_out}, input.options());\n\n auto grad_input = at::zeros_like(input);\n auto grad_weight = at::zeros_like(weight);\n auto grad_bias = at::zeros_like(bias);\n auto grad_offset = at::zeros_like(offset);\n auto grad_mask = at::zeros_like(mask);\n\n using scalar_t = float;\n\n for (int b = 0; b < batch; b++)\n {\n auto input_n = input.select(0, b);\n auto offset_n = offset.select(0, b);\n auto mask_n = mask.select(0, b);\n auto grad_output_n = grad_output.select(0, b);\n auto grad_input_n = grad_input.select(0, b);\n auto grad_offset_n = grad_offset.select(0, b);\n auto grad_mask_n = grad_mask.select(0, b);\n\n long m = channels * kernel_h * kernel_w;\n long n = height_out * width_out;\n long k = channels_out;\n\n THFloatBlas_gemm('n', 't', n, m, k, 1.0f,\n grad_output_n.data(), n,\n weight.data(), m, 0.0f,\n columns.data(), n);\n\n // gradient w.r.t. input coordinate data\n modulated_deformable_col2im_coord_cpu(columns.data(),\n input_n.data(),\n offset_n.data(),\n mask_n.data(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, deformable_group,\n grad_offset_n.data(),\n grad_mask_n.data());\n // gradient w.r.t. input data\n modulated_deformable_col2im_cpu(columns.data(),\n offset_n.data(),\n mask_n.data(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, deformable_group,\n grad_input_n.data());\n\n // gradient w.r.t. weight, dWeight should accumulate across the batch and group\n modulated_deformable_im2col_cpu(input_n.data(),\n offset_n.data(),\n mask_n.data(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, deformable_group,\n columns.data());\n\n long m_ = channels_out;\n long n_ = channels * kernel_h * kernel_w;\n long k_ = height_out * width_out;\n\n THFloatBlas_gemm('t', 'n', n_, m_, k_, 1.0f,\n columns.data(), k_,\n grad_output_n.data(), k_, 1.0f,\n grad_weight.data(), n_);\n\n // gradient w.r.t. bias\n // long m_ = channels_out;\n // long k__ = height_out * width_out;\n // THFloatBlas_gemv('t', k_, m_, 1.0f,\n // grad_output_n.data(), k_,\n // ones.data(), 1, 1.0f,\n // grad_bias.data(), 1);\n }\n\n return {\n grad_input, grad_offset, grad_mask, grad_weight, grad_bias\n };\n}"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cpu/dcn_v2_im2col_cpu.cpp",
"content": "#include \"dcn_v2_im2col_cpu.h\"\n#include \n#include \n#include \n\n#include \n//#include \n\n#include \n//#include \n//#include \n\n// modified from the CUDA version for CPU use by Daniel K. Suhendro\n\n/*#define CUDA_KERNEL_LOOP(i, n) \\\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; \\\n i < (n); \\\n i += blockDim.x * gridDim.x)\n\nconst int CUDA_NUM_THREADS = 1024;\ninline int GET_BLOCKS(const int N)\n{\n return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;\n}*/\n\n\nfloat dmcn_im2col_bilinear_cpu(const float *bottom_data, const int data_width,\n const int height, const int width, float h, float w)\n{\n int h_low = floor(h);\n int w_low = floor(w);\n int h_high = h_low + 1;\n int w_high = w_low + 1;\n\n float lh = h - h_low;\n float lw = w - w_low;\n float hh = 1 - lh, hw = 1 - lw;\n\n float v1 = 0;\n if (h_low >= 0 && w_low >= 0)\n v1 = bottom_data[h_low * data_width + w_low];\n float v2 = 0;\n if (h_low >= 0 && w_high <= width - 1)\n v2 = bottom_data[h_low * data_width + w_high];\n float v3 = 0;\n if (h_high <= height - 1 && w_low >= 0)\n v3 = bottom_data[h_high * data_width + w_low];\n float v4 = 0;\n if (h_high <= height - 1 && w_high <= width - 1)\n v4 = bottom_data[h_high * data_width + w_high];\n\n float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;\n\n float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);\n return val;\n}\n\nfloat dmcn_get_gradient_weight_cpu(float argmax_h, float argmax_w,\n const int h, const int w, const int height, const int width)\n{\n if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)\n {\n //empty\n return 0;\n }\n\n int argmax_h_low = floor(argmax_h);\n int argmax_w_low = floor(argmax_w);\n int argmax_h_high = argmax_h_low + 1;\n int argmax_w_high = argmax_w_low + 1;\n\n float weight = 0;\n if (h == argmax_h_low && w == argmax_w_low)\n weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);\n if (h == argmax_h_low && w == argmax_w_high)\n weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);\n if (h == argmax_h_high && w == argmax_w_low)\n weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);\n if (h == argmax_h_high && w == argmax_w_high)\n weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);\n return weight;\n}\n\nfloat dmcn_get_coordinate_weight_cpu(float argmax_h, float argmax_w,\n const int height, const int width, const float *im_data,\n const int data_width, const int bp_dir)\n{\n if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)\n {\n //empty\n return 0;\n }\n\n int argmax_h_low = floor(argmax_h);\n int argmax_w_low = floor(argmax_w);\n int argmax_h_high = argmax_h_low + 1;\n int argmax_w_high = argmax_w_low + 1;\n\n float weight = 0;\n\n if (bp_dir == 0)\n {\n if (argmax_h_low >= 0 && argmax_w_low >= 0)\n weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];\n if (argmax_h_low >= 0 && argmax_w_high <= width - 1)\n weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];\n if (argmax_h_high <= height - 1 && argmax_w_low >= 0)\n weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];\n if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)\n weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];\n }\n else if (bp_dir == 1)\n {\n if (argmax_h_low >= 0 && argmax_w_low >= 0)\n weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];\n if (argmax_h_low >= 0 && argmax_w_high <= width - 1)\n weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];\n if (argmax_h_high <= height - 1 && argmax_w_low >= 0)\n weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];\n if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)\n weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];\n }\n\n return weight;\n}\n\nvoid modulated_deformable_im2col_cpu_kernel(const int n, const float *data_im, const float *data_offset, const float *data_mask,\n const int height, const int width, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w,\n const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int channel_per_deformable_group,\n const int batch_size, const int num_channels, const int deformable_group,\n const int height_col, const int width_col,\n float *data_col)\n{\n // launch channels * batch_size * height_col * width_col cores\n for(int index=0; index(0);\n const float h_im = h_in + i * dilation_h + offset_h;\n const float w_im = w_in + j * dilation_w + offset_w;\n //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {\n if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)\n {\n //const float map_h = i * dilation_h + offset_h;\n //const float map_w = j * dilation_w + offset_w;\n //const int cur_height = height - h_in;\n //const int cur_width = width - w_in;\n //val = dmcn_im2col_bilinear_cpu(data_im_ptr, width, cur_height, cur_width, map_h, map_w);\n val = dmcn_im2col_bilinear_cpu(data_im_ptr, width, height, width, h_im, w_im);\n }\n *data_col_ptr = val * mask;\n // data_col_ptr += batch_size * height_col * width_col;\n data_col_ptr += height_col * width_col;\n }\n }\n }\n}\n\nvoid modulated_deformable_col2im_cpu_kernel(const int n, const float *data_col, const float *data_offset, const float *data_mask,\n const int channels, const int height, const int width,\n const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w,\n const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int channel_per_deformable_group,\n const int batch_size, const int deformable_group,\n const int height_col, const int width_col,\n float *grad_im)\n{\n for(int index = 0; index < n; index++)\n {\n const int j = (index / width_col / height_col / batch_size) % kernel_w;\n const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;\n const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;\n // compute the start and end of the output\n\n const int deformable_group_index = c / channel_per_deformable_group;\n\n int w_out = index % width_col;\n int h_out = (index / width_col) % height_col;\n int b = (index / width_col / height_col) % batch_size;\n int w_in = w_out * stride_w - pad_w;\n int h_in = h_out * stride_h - pad_h;\n\n const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;\n const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;\n const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;\n const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;\n const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_out) * width_col + w_out;\n const float offset_h = data_offset_ptr[data_offset_h_ptr];\n const float offset_w = data_offset_ptr[data_offset_w_ptr];\n const float mask = data_mask_ptr[data_mask_hw_ptr];\n const float cur_inv_h_data = h_in + i * dilation_h + offset_h;\n const float cur_inv_w_data = w_in + j * dilation_w + offset_w;\n\n const float cur_top_grad = data_col[index] * mask;\n const int cur_h = (int)cur_inv_h_data;\n const int cur_w = (int)cur_inv_w_data;\n \n for (int dy = -2; dy <= 2; dy++)\n {\n for (int dx = -2; dx <= 2; dx++)\n {\n if (cur_h + dy >= 0 && cur_h + dy < height &&\n cur_w + dx >= 0 && cur_w + dx < width &&\n abs(cur_inv_h_data - (cur_h + dy)) < 1 &&\n abs(cur_inv_w_data - (cur_w + dx)) < 1)\n {\n int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;\n float weight = dmcn_get_gradient_weight_cpu(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);\n //atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);\n *(grad_im + cur_bottom_grad_pos) += weight * cur_top_grad;\n\n }\n }\n }\n }\n}\n\nvoid modulated_deformable_col2im_coord_cpu_kernel(const int n, const float *data_col, const float *data_im,\n const float *data_offset, const float *data_mask,\n const int channels, const int height, const int width,\n const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w,\n const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int channel_per_deformable_group,\n const int batch_size, const int offset_channels, const int deformable_group,\n const int height_col, const int width_col,\n float *grad_offset, float *grad_mask)\n{\n for(int index = 0; index < n; index++)\n {\n float val = 0, mval = 0;\n int w = index % width_col;\n int h = (index / width_col) % height_col;\n int c = (index / width_col / height_col) % offset_channels;\n int b = (index / width_col / height_col) / offset_channels;\n // compute the start and end of the output\n\n const int deformable_group_index = c / (2 * kernel_h * kernel_w);\n const int col_step = kernel_h * kernel_w;\n int cnt = 0;\n const float *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group * batch_size * width_col * height_col;\n const float *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) * channel_per_deformable_group / kernel_h / kernel_w * height * width;\n const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;\n const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;\n\n const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;\n\n for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)\n {\n const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;\n const int bp_dir = offset_c % 2;\n\n int j = (col_pos / width_col / height_col / batch_size) % kernel_w;\n int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;\n int w_out = col_pos % width_col;\n int h_out = (col_pos / width_col) % height_col;\n int w_in = w_out * stride_w - pad_w;\n int h_in = h_out * stride_h - pad_h;\n const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);\n const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);\n const int data_mask_hw_ptr = (((i * kernel_w + j) * height_col + h_out) * width_col + w_out);\n const float offset_h = data_offset_ptr[data_offset_h_ptr];\n const float offset_w = data_offset_ptr[data_offset_w_ptr];\n const float mask = data_mask_ptr[data_mask_hw_ptr];\n float inv_h = h_in + i * dilation_h + offset_h;\n float inv_w = w_in + j * dilation_w + offset_w;\n if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)\n {\n inv_h = inv_w = -2;\n }\n else\n {\n mval += data_col_ptr[col_pos] * dmcn_im2col_bilinear_cpu(data_im_ptr + cnt * height * width, width, height, width, inv_h, inv_w);\n }\n const float weight = dmcn_get_coordinate_weight_cpu(\n inv_h, inv_w,\n height, width, data_im_ptr + cnt * height * width, width, bp_dir);\n val += weight * data_col_ptr[col_pos] * mask;\n cnt += 1;\n }\n // KERNEL_ASSIGN(grad_offset[index], offset_req, val);\n grad_offset[index] = val;\n if (offset_c % 2 == 0)\n // KERNEL_ASSIGN(grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w], mask_req, mval);\n grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w] = mval;\n }\n}\n\nvoid modulated_deformable_im2col_cpu(const float* data_im, const float* data_offset, const float* data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im, \n const int height_col, const int width_col, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w, \n const int dilation_h, const int dilation_w,\n const int deformable_group, float* data_col) {\n // num_axes should be smaller than block size\n const int channel_per_deformable_group = channels / deformable_group;\n const int num_kernels = channels * batch_size * height_col * width_col;\n modulated_deformable_im2col_cpu_kernel(\n num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group,\n batch_size, channels, deformable_group, height_col, width_col, data_col);\n \n /*cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in modulated_deformable_im2col_cuda: %s\\n\", cudaGetErrorString(err));\n }*/\n\n}\n\nvoid modulated_deformable_col2im_cpu(const float* data_col, const float* data_offset, const float* data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im, \n const int height_col, const int width_col, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w, \n const int dilation_h, const int dilation_w, \n const int deformable_group, float* grad_im){\n\n const int channel_per_deformable_group = channels / deformable_group;\n const int num_kernels = channels * kernel_h * kernel_w * batch_size * height_col * width_col;\n modulated_deformable_col2im_cpu_kernel(\n num_kernels, data_col, data_offset, data_mask, channels, height_im, width_im,\n kernel_h, kernel_w, pad_h, pad_h, stride_h, stride_w,\n dilation_h, dilation_w, channel_per_deformable_group,\n batch_size, deformable_group, height_col, width_col, grad_im);\n /*cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in modulated_deformable_col2im_cuda: %s\\n\", cudaGetErrorString(err));\n }*/\n\n}\n\nvoid modulated_deformable_col2im_coord_cpu(const float* data_col, const float* data_im, const float* data_offset, const float* data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im, \n const int height_col, const int width_col, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w, \n const int dilation_h, const int dilation_w, \n const int deformable_group,\n float* grad_offset, float* grad_mask) {\n const int num_kernels = batch_size * height_col * width_col * 2 * kernel_h * kernel_w * deformable_group;\n const int channel_per_deformable_group = channels * kernel_h * kernel_w / deformable_group;\n modulated_deformable_col2im_coord_cpu_kernel(\n num_kernels, data_col, data_im, data_offset, data_mask, channels, height_im, width_im,\n kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, channel_per_deformable_group,\n batch_size, 2 * kernel_h * kernel_w * deformable_group, deformable_group, height_col, width_col, \n grad_offset, grad_mask);\n /*cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in modulated_deformable_col2im_coord_cuda: %s\\n\", cudaGetErrorString(err));\n }*/\n}"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cpu/dcn_v2_im2col_cpu.h",
"content": "\n/*!\n ******************* BEGIN Caffe Copyright Notice and Disclaimer ****************\n *\n * COPYRIGHT\n *\n * All contributions by the University of California:\n * Copyright (c) 2014-2017 The Regents of the University of California (Regents)\n * All rights reserved.\n *\n * All other contributions:\n * Copyright (c) 2014-2017, the respective contributors\n * All rights reserved.\n *\n * Caffe uses a shared copyright model: each contributor holds copyright over\n * their contributions to Caffe. The project versioning records all such\n * contribution and copyright details. If a contributor wants to further mark\n * their specific copyright on a particular contribution, they should indicate\n * their copyright solely in the commit message of the change when it is\n * committed.\n *\n * LICENSE\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n *\n * 1. Redistributions of source code must retain the above copyright notice, this\n * list of conditions and the following disclaimer.\n * 2. Redistributions in binary form must reproduce the above copyright notice,\n * this list of conditions and the following disclaimer in the documentation\n * and/or other materials provided with the distribution.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR\n * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n * CONTRIBUTION AGREEMENT\n *\n * By contributing to the BVLC/caffe repository through pull-request, comment,\n * or otherwise, the contributor releases their content to the\n * license and copyright terms herein.\n *\n ***************** END Caffe Copyright Notice and Disclaimer ********************\n *\n * Copyright (c) 2018 Microsoft\n * Licensed under The MIT License [see LICENSE for details]\n * \\file modulated_deformable_im2col.h\n * \\brief Function definitions of converting an image to\n * column matrix based on kernel, padding, dilation, and offset.\n * These functions are mainly used in deformable convolution operators.\n * \\ref: https://arxiv.org/abs/1811.11168\n * \\author Yuwen Xiong, Haozhi Qi, Jifeng Dai, Xizhou Zhu, Han Hu\n */\n\n/***************** Adapted by Charles Shang *********************/\n// modified from the CUDA version for CPU use by Daniel K. Suhendro\n\n#ifndef DCN_V2_IM2COL_CPU\n#define DCN_V2_IM2COL_CPU\n\n#ifdef __cplusplus\nextern \"C\"\n{\n#endif\n\n void modulated_deformable_im2col_cpu(const float *data_im, const float *data_offset, const float *data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im,\n const int height_col, const int width_col, const int kernel_h, const int kenerl_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int deformable_group, float *data_col);\n\n void modulated_deformable_col2im_cpu(const float *data_col, const float *data_offset, const float *data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im,\n const int height_col, const int width_col, const int kernel_h, const int kenerl_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int deformable_group, float *grad_im);\n\n void modulated_deformable_col2im_coord_cpu(const float *data_col, const float *data_im, const float *data_offset, const float *data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im,\n const int height_col, const int width_col, const int kernel_h, const int kenerl_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int deformable_group,\n float *grad_offset, float *grad_mask);\n\n#ifdef __cplusplus\n}\n#endif\n\n#endif"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cpu/dcn_v2_psroi_pooling_cpu.cpp",
"content": "/*!\n * Copyright (c) 2017 Microsoft\n * Licensed under The MIT License [see LICENSE for details]\n * \\file deformable_psroi_pooling.cu\n * \\brief\n * \\author Yi Li, Guodong Zhang, Jifeng Dai\n*/\n/***************** Adapted by Charles Shang *********************/\n// modified from the CUDA version for CPU use by Daniel K. Suhendro\n\n#include \n#include \n#include \n\n#include \n//#include \n\n#include | \n//#include \n//#include \n\n/*#define CUDA_KERNEL_LOOP(i, n) \\\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; \\\n i < (n); \\\n i += blockDim.x * gridDim.x)\n\nconst int CUDA_NUM_THREADS = 1024;\ninline int GET_BLOCKS(const int N)\n{\n return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;\n}*/\n\ntemplate \nT bilinear_interp_cpu(\n const T *data,\n const T x,\n const T y,\n const int width,\n const int height)\n{\n int x1 = floor(x);\n int x2 = ceil(x);\n int y1 = floor(y);\n int y2 = ceil(y);\n T dist_x = static_cast(x - x1);\n T dist_y = static_cast(y - y1);\n T value11 = data[y1 * width + x1];\n T value12 = data[y2 * width + x1];\n T value21 = data[y1 * width + x2];\n T value22 = data[y2 * width + x2];\n T value = (1 - dist_x) * (1 - dist_y) * value11 +\n (1 - dist_x) * dist_y * value12 +\n dist_x * (1 - dist_y) * value21 +\n dist_x * dist_y * value22;\n return value;\n}\n\ntemplate \n void DeformablePSROIPoolForwardKernelCpu(\n const int count,\n const T *bottom_data,\n const T spatial_scale,\n const int channels,\n const int height, const int width,\n const int pooled_height, const int pooled_width,\n const T *bottom_rois, const T *bottom_trans,\n const int no_trans,\n const T trans_std,\n const int sample_per_part,\n const int output_dim,\n const int group_size,\n const int part_size,\n const int num_classes,\n const int channels_each_class,\n T *top_data,\n T *top_count)\n{\n for(int index = 0; index < count; index++)\n {\n // The output is in order (n, ctop, ph, pw)\n int pw = index % pooled_width;\n int ph = (index / pooled_width) % pooled_height;\n int ctop = (index / pooled_width / pooled_height) % output_dim;\n int n = index / pooled_width / pooled_height / output_dim;\n\n // [start, end) interval for spatial sampling\n const T *offset_bottom_rois = bottom_rois + n * 5;\n int roi_batch_ind = offset_bottom_rois[0];\n T roi_start_w = static_cast(round(offset_bottom_rois[1])) * spatial_scale - 0.5;\n T roi_start_h = static_cast(round(offset_bottom_rois[2])) * spatial_scale - 0.5;\n T roi_end_w = static_cast(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;\n T roi_end_h = static_cast(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;\n\n // Force too small ROIs to be 1x1\n T roi_width = std::max(roi_end_w - roi_start_w, T(0.1)); //avoid 0\n T roi_height = std::max(roi_end_h - roi_start_h, T(0.1));\n\n // Compute w and h at bottom\n T bin_size_h = roi_height / static_cast(pooled_height);\n T bin_size_w = roi_width / static_cast(pooled_width);\n\n T sub_bin_size_h = bin_size_h / static_cast(sample_per_part);\n T sub_bin_size_w = bin_size_w / static_cast(sample_per_part);\n\n int part_h = floor(static_cast(ph) / pooled_height * part_size);\n int part_w = floor(static_cast(pw) / pooled_width * part_size);\n int class_id = ctop / channels_each_class;\n T trans_x = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * trans_std;\n T trans_y = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * trans_std;\n\n T wstart = static_cast(pw) * bin_size_w + roi_start_w;\n wstart += trans_x * roi_width;\n T hstart = static_cast(ph) * bin_size_h + roi_start_h;\n hstart += trans_y * roi_height;\n\n T sum = 0;\n int count = 0;\n int gw = floor(static_cast(pw) * group_size / pooled_width);\n int gh = floor(static_cast(ph) * group_size / pooled_height);\n gw = std::min(std::max(gw, 0), group_size - 1);\n gh = std::min(std::max(gh, 0), group_size - 1);\n\n const T *offset_bottom_data = bottom_data + (roi_batch_ind * channels) * height * width;\n for (int ih = 0; ih < sample_per_part; ih++)\n {\n for (int iw = 0; iw < sample_per_part; iw++)\n {\n T w = wstart + iw * sub_bin_size_w;\n T h = hstart + ih * sub_bin_size_h;\n // bilinear interpolation\n if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)\n {\n continue;\n }\n w = std::min(std::max(w, T(0.)), width - T(1.));\n h = std::min(std::max(h, T(0.)), height - T(1.));\n int c = (ctop * group_size + gh) * group_size + gw;\n T val = bilinear_interp_cpu(offset_bottom_data + c * height * width, w, h, width, height);\n sum += val;\n count++;\n }\n }\n top_data[index] = count == 0 ? static_cast(0) : sum / count;\n top_count[index] = count;\n }\n}\n\ntemplate \nvoid DeformablePSROIPoolBackwardAccKernelCpu(\n const int count,\n const T *top_diff,\n const T *top_count,\n const int num_rois,\n const T spatial_scale,\n const int channels,\n const int height, const int width,\n const int pooled_height, const int pooled_width,\n const int output_dim,\n T *bottom_data_diff, T *bottom_trans_diff,\n const T *bottom_data,\n const T *bottom_rois,\n const T *bottom_trans,\n const int no_trans,\n const T trans_std,\n const int sample_per_part,\n const int group_size,\n const int part_size,\n const int num_classes,\n const int channels_each_class)\n{\n for(int index = 0; index < count; index++)\n {\n // The output is in order (n, ctop, ph, pw)\n int pw = index % pooled_width;\n int ph = (index / pooled_width) % pooled_height;\n int ctop = (index / pooled_width / pooled_height) % output_dim;\n int n = index / pooled_width / pooled_height / output_dim;\n\n // [start, end) interval for spatial sampling\n const T *offset_bottom_rois = bottom_rois + n * 5;\n int roi_batch_ind = offset_bottom_rois[0];\n T roi_start_w = static_cast(round(offset_bottom_rois[1])) * spatial_scale - 0.5;\n T roi_start_h = static_cast(round(offset_bottom_rois[2])) * spatial_scale - 0.5;\n T roi_end_w = static_cast(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;\n T roi_end_h = static_cast(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;\n \n // Force too small ROIs to be 1x1\n T roi_width = std::max(roi_end_w - roi_start_w, T(0.1)); //avoid 0\n T roi_height = std::max(roi_end_h - roi_start_h, T(0.1));\n\n // Compute w and h at bottom\n T bin_size_h = roi_height / static_cast(pooled_height);\n T bin_size_w = roi_width / static_cast(pooled_width);\n\n T sub_bin_size_h = bin_size_h / static_cast(sample_per_part);\n T sub_bin_size_w = bin_size_w / static_cast(sample_per_part);\n\n int part_h = floor(static_cast(ph) / pooled_height * part_size);\n int part_w = floor(static_cast(pw) / pooled_width * part_size);\n int class_id = ctop / channels_each_class;\n T trans_x = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * trans_std;\n T trans_y = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * trans_std;\n\n T wstart = static_cast(pw) * bin_size_w + roi_start_w;\n wstart += trans_x * roi_width;\n T hstart = static_cast(ph) * bin_size_h + roi_start_h;\n hstart += trans_y * roi_height;\n\n if (top_count[index] <= 0)\n {\n continue;\n }\n T diff_val = top_diff[index] / top_count[index];\n const T *offset_bottom_data = bottom_data + roi_batch_ind * channels * height * width;\n T *offset_bottom_data_diff = bottom_data_diff + roi_batch_ind * channels * height * width;\n int gw = floor(static_cast(pw) * group_size / pooled_width);\n int gh = floor(static_cast(ph) * group_size / pooled_height);\n gw = std::min(std::max(gw, 0), group_size - 1);\n gh = std::min(std::max(gh, 0), group_size - 1);\n\n for (int ih = 0; ih < sample_per_part; ih++)\n {\n for (int iw = 0; iw < sample_per_part; iw++)\n {\n T w = wstart + iw * sub_bin_size_w;\n T h = hstart + ih * sub_bin_size_h;\n // bilinear interpolation\n if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)\n {\n continue;\n }\n w = std::min(std::max(w, T(0.)), width - T(1.));\n h = std::min(std::max(h, T(0.)), height - T(1.));\n int c = (ctop * group_size + gh) * group_size + gw;\n // backward on feature\n int x0 = floor(w);\n int x1 = ceil(w);\n int y0 = floor(h);\n int y1 = ceil(h);\n T dist_x = w - x0, dist_y = h - y0;\n T q00 = (1 - dist_x) * (1 - dist_y);\n T q01 = (1 - dist_x) * dist_y;\n T q10 = dist_x * (1 - dist_y);\n T q11 = dist_x * dist_y;\n int bottom_index_base = c * height * width;\n /*atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x0, q00 * diff_val);\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x0, q01 * diff_val);\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x1, q10 * diff_val);\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x1, q11 * diff_val);*/\n *(offset_bottom_data_diff + bottom_index_base + y0 * width + x0) += q00 * diff_val;\n *(offset_bottom_data_diff + bottom_index_base + y1 * width + x0) += q01 * diff_val;\n *(offset_bottom_data_diff + bottom_index_base + y0 * width + x1) += q10 * diff_val;\n *(offset_bottom_data_diff + bottom_index_base + y1 * width + x1) += q11 * diff_val;\n\n\n if (no_trans)\n {\n continue;\n }\n T U00 = offset_bottom_data[bottom_index_base + y0 * width + x0];\n T U01 = offset_bottom_data[bottom_index_base + y1 * width + x0];\n T U10 = offset_bottom_data[bottom_index_base + y0 * width + x1];\n T U11 = offset_bottom_data[bottom_index_base + y1 * width + x1];\n T diff_x = (U11 * dist_y + U10 * (1 - dist_y) - U01 * dist_y - U00 * (1 - dist_y)) * trans_std * diff_val;\n diff_x *= roi_width;\n T diff_y = (U11 * dist_x + U01 * (1 - dist_x) - U10 * dist_x - U00 * (1 - dist_x)) * trans_std * diff_val;\n diff_y *= roi_height;\n\n /*atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w, diff_x);\n atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w, diff_y);*/\n *(bottom_trans_diff + (((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w) += diff_x;\n *(bottom_trans_diff + (((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w) += diff_y;\n }\n }\n }\n}\n\nstd::tuple\ndcn_v2_psroi_pooling_cpu_forward(const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std)\n{\n /*AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(bbox.type().is_cuda(), \"rois must be a CUDA tensor\");\n AT_ASSERTM(trans.type().is_cuda(), \"trans must be a CUDA tensor\");*/\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n const int channels_trans = no_trans ? 2 : trans.size(1);\n const int num_bbox = bbox.size(0);\n\n AT_ASSERTM(channels == output_dim, \"input channels and output channels must equal\");\n auto pooled_height = pooled_size;\n auto pooled_width = pooled_size;\n\n auto out = at::empty({num_bbox, output_dim, pooled_height, pooled_width}, input.options());\n long out_size = num_bbox * output_dim * pooled_height * pooled_width;\n auto top_count = at::zeros({num_bbox, output_dim, pooled_height, pooled_width}, input.options());\n\n const int num_classes = no_trans ? 1 : channels_trans / 2;\n const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;\n\n //cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (out.numel() == 0)\n {\n //THCudaCheck(cudaGetLastError());\n return std::make_tuple(out, top_count);\n }\n\n /*dim3 grid(std::min(THCCeilDiv(out_size, 512L), 4096L));\n dim3 block(512);*/\n\n AT_DISPATCH_FLOATING_TYPES(input.type(), \"dcn_v2_psroi_pooling_cpu_forward\", [&] {\n DeformablePSROIPoolForwardKernelCpu(\n out_size,\n input.contiguous().data(),\n spatial_scale,\n channels,\n height, width,\n pooled_height,\n pooled_width,\n bbox.contiguous().data(),\n trans.contiguous().data(),\n no_trans,\n trans_std,\n sample_per_part,\n output_dim,\n group_size,\n part_size,\n num_classes,\n channels_each_class,\n out.data(),\n top_count.data());\n });\n //THCudaCheck(cudaGetLastError());\n return std::make_tuple(out, top_count);\n}\n\nstd::tuple\ndcn_v2_psroi_pooling_cpu_backward(const at::Tensor &out_grad,\n const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const at::Tensor &top_count,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std)\n{\n /*AT_ASSERTM(out_grad.type().is_cuda(), \"out_grad must be a CUDA tensor\");\n AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(bbox.type().is_cuda(), \"bbox must be a CUDA tensor\");\n AT_ASSERTM(trans.type().is_cuda(), \"trans must be a CUDA tensor\");\n AT_ASSERTM(top_count.type().is_cuda(), \"top_count must be a CUDA tensor\");*/\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n const int channels_trans = no_trans ? 2 : trans.size(1);\n const int num_bbox = bbox.size(0);\n\n AT_ASSERTM(channels == output_dim, \"input channels and output channels must equal\");\n auto pooled_height = pooled_size;\n auto pooled_width = pooled_size;\n long out_size = num_bbox * output_dim * pooled_height * pooled_width;\n const int num_classes = no_trans ? 1 : channels_trans / 2;\n const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;\n\n auto input_grad = at::zeros({batch, channels, height, width}, out_grad.options());\n auto trans_grad = at::zeros_like(trans);\n\n if (input_grad.numel() == 0)\n {\n //THCudaCheck(cudaGetLastError());\n return std::make_tuple(input_grad, trans_grad);\n }\n\n /*dim3 grid(std::min(THCCeilDiv(out_size, 512L), 4096L));\n dim3 block(512);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();*/\n\n AT_DISPATCH_FLOATING_TYPES(out_grad.type(), \"dcn_v2_psroi_pooling_cpu_backward\", [&] {\n DeformablePSROIPoolBackwardAccKernelCpu(\n out_size,\n out_grad.contiguous().data(),\n top_count.contiguous().data(),\n num_bbox,\n spatial_scale,\n channels,\n height,\n width,\n pooled_height,\n pooled_width,\n output_dim,\n input_grad.contiguous().data(),\n trans_grad.contiguous().data(),\n input.contiguous().data(),\n bbox.contiguous().data(),\n trans.contiguous().data(),\n no_trans,\n trans_std,\n sample_per_part,\n group_size,\n part_size,\n num_classes,\n channels_each_class);\n });\n //THCudaCheck(cudaGetLastError());\n return std::make_tuple(input_grad, trans_grad);\n}"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cpu/vision.h",
"content": "#pragma once\n#include \n\nat::Tensor\ndcn_v2_cpu_forward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const int kernel_h,\n const int kernel_w,\n const int stride_h,\n const int stride_w,\n const int pad_h,\n const int pad_w,\n const int dilation_h,\n const int dilation_w,\n const int deformable_group);\n\nstd::vector\ndcn_v2_cpu_backward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const at::Tensor &grad_output,\n int kernel_h, int kernel_w,\n int stride_h, int stride_w,\n int pad_h, int pad_w,\n int dilation_h, int dilation_w,\n int deformable_group);\n\n\nstd::tuple\ndcn_v2_psroi_pooling_cpu_forward(const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std);\n\nstd::tuple\ndcn_v2_psroi_pooling_cpu_backward(const at::Tensor &out_grad,\n const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const at::Tensor &top_count,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std);"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cuda/dcn_v2_cuda.cu",
"content": "#include \n#include \"cuda/dcn_v2_im2col_cuda.h\"\n\n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n#include \n\nTHCState *state = at::globalContext().lazyInitCUDA();\n\nstatic cublasOperation_t _cublasOpFromChar(char op) {\n switch (op) {\n case 'n':\n case 'N':\n return CUBLAS_OP_N;\n case 't':\n case 'T':\n return CUBLAS_OP_T;\n case 'c':\n case 'C':\n return CUBLAS_OP_C;\n }\n AT_ERROR(\n \"_cublasOpFromChar input should be 't', 'n' or 'c' but got `\", op, \"`\");\n }\n\n static void _cublasAdjustLdLevel2(int64_t m, int64_t n, int64_t* lda) {\n // Note: leading dimensions generally are checked that they are > 0\n // and at least as big the result requires (even if the value won't\n // be used).\n \n // Q: Why does Level3 check trans but this doesn't?\n // A: In level 2, the sizes (m, n) specify the size of A\n // (independent of trans value). In level 3. the sizes (m, n, k)\n // specify the sizes of op(A), op(B) where op depend on trans\n // values.\n if (n <= 1)\n *lda = std::max(m, 1);\n }\n\n\n\n// author: Charles Shang\n// https://github.com/torch/cunn/blob/master/lib/THCUNN/generic/SpatialConvolutionMM.cu\n\n// [batch gemm]\n// https://github.com/pytorch/pytorch/blob/master/aten/src/THC/generic/THCTensorMathBlas.cu\n\n__global__ void createBatchGemmBuffer(const float **input_b, float **output_b,\n float **columns_b, const float **ones_b,\n const float **weight_b, const float **bias_b,\n float *input, float *output,\n float *columns, float *ones,\n float *weight, float *bias,\n const int input_stride, const int output_stride,\n const int columns_stride, const int ones_stride,\n const int num_batches)\n{\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < num_batches)\n {\n input_b[idx] = input + idx * input_stride;\n output_b[idx] = output + idx * output_stride;\n columns_b[idx] = columns + idx * columns_stride;\n ones_b[idx] = ones + idx * ones_stride;\n // share weights and bias within a Mini-Batch\n weight_b[idx] = weight;\n bias_b[idx] = bias;\n }\n}\n\nat::Tensor\ndcn_v2_cuda_forward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const int kernel_h,\n const int kernel_w,\n const int stride_h,\n const int stride_w,\n const int pad_h,\n const int pad_w,\n const int dilation_h,\n const int dilation_w,\n const int deformable_group)\n{\n using scalar_t = float;\n // THCAssertSameGPU(THCudaTensor_checkGPU(state, 5, input, weight, bias, offset, mask));\n AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(weight.type().is_cuda(), \"weight must be a CUDA tensor\");\n AT_ASSERTM(bias.type().is_cuda(), \"bias must be a CUDA tensor\");\n AT_ASSERTM(offset.type().is_cuda(), \"offset must be a CUDA tensor\");\n AT_ASSERTM(mask.type().is_cuda(), \"mask must be a CUDA tensor\");\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n\n const int channels_out = weight.size(0);\n const int channels_kernel = weight.size(1);\n const int kernel_h_ = weight.size(2);\n const int kernel_w_ = weight.size(3);\n\n // printf(\"Kernels: %d %d %d %d\\n\", kernel_h_, kernel_w_, kernel_w, kernel_h);\n // printf(\"Channels: %d %d\\n\", channels, channels_kernel);\n // printf(\"Channels: %d %d\\n\", channels_out, channels_kernel);\n\n AT_ASSERTM(kernel_h_ == kernel_h && kernel_w_ == kernel_w,\n \"Input shape and kernel shape wont match: (%d x %d vs %d x %d).\", kernel_h_, kernel_w, kernel_h_, kernel_w_);\n\n AT_ASSERTM(channels == channels_kernel,\n \"Input shape and kernel channels wont match: (%d vs %d).\", channels, channels_kernel);\n\n const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;\n const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;\n\n auto ones = at::ones({batch, height_out, width_out}, input.options());\n auto columns = at::empty({batch, channels * kernel_h * kernel_w, 1 * height_out * width_out}, input.options());\n auto output = at::empty({batch, channels_out, height_out, width_out}, input.options());\n\n // prepare for batch-wise computing, which is significantly faster than instance-wise computing\n // when batch size is large.\n // launch batch threads\n int matrices_size = batch * sizeof(float *);\n auto input_b = static_cast(THCudaMalloc(state, matrices_size));\n auto output_b = static_cast(THCudaMalloc(state, matrices_size));\n auto columns_b = static_cast(THCudaMalloc(state, matrices_size));\n auto ones_b = static_cast(THCudaMalloc(state, matrices_size));\n auto weight_b = static_cast(THCudaMalloc(state, matrices_size));\n auto bias_b = static_cast(THCudaMalloc(state, matrices_size));\n\n const int block = 128;\n const int grid = (batch + block - 1) / block;\n\n createBatchGemmBuffer<<>>(\n input_b, output_b,\n columns_b, ones_b,\n weight_b, bias_b,\n input.data_ptr(),\n output.data_ptr(),\n columns.data_ptr(),\n ones.data_ptr(),\n weight.data_ptr(),\n bias.data_ptr(),\n channels * width * height,\n channels_out * width_out * height_out,\n channels * kernel_h * kernel_w * height_out * width_out,\n height_out * width_out,\n batch);\n\n long m_ = channels_out;\n long n_ = height_out * width_out;\n long k_ = 1;\n THCudaBlas_SgemmBatched(state,\n 't',\n 'n',\n n_,\n m_,\n k_,\n 1.0f,\n ones_b, k_,\n bias_b, k_,\n 0.0f,\n output_b, n_,\n batch);\n\n modulated_deformable_im2col_cuda(c10::cuda::getCurrentCUDAStream(),\n input.data_ptr(),\n offset.data_ptr(),\n mask.data_ptr(),\n batch, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w,\n deformable_group,\n columns.data_ptr());\n\n long m = channels_out;\n long n = height_out * width_out;\n long k = channels * kernel_h * kernel_w;\n THCudaBlas_SgemmBatched(state,\n 'n',\n 'n',\n n,\n m,\n k,\n 1.0f,\n (const float **)columns_b, n,\n weight_b, k,\n 1.0f,\n output_b, n,\n batch);\n\n THCudaFree(state, input_b);\n THCudaFree(state, output_b);\n THCudaFree(state, columns_b);\n THCudaFree(state, ones_b);\n THCudaFree(state, weight_b);\n THCudaFree(state, bias_b);\n return output;\n}\n\n__global__ void createBatchGemmBufferBackward(\n float **grad_output_b,\n float **columns_b,\n float **ones_b,\n float **weight_b,\n float **grad_weight_b,\n float **grad_bias_b,\n float *grad_output,\n float *columns,\n float *ones,\n float *weight,\n float *grad_weight,\n float *grad_bias,\n const int grad_output_stride,\n const int columns_stride,\n const int ones_stride,\n const int num_batches)\n{\n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n if (idx < num_batches)\n {\n grad_output_b[idx] = grad_output + idx * grad_output_stride;\n columns_b[idx] = columns + idx * columns_stride;\n ones_b[idx] = ones + idx * ones_stride;\n\n // share weights and bias within a Mini-Batch\n weight_b[idx] = weight;\n grad_weight_b[idx] = grad_weight;\n grad_bias_b[idx] = grad_bias;\n }\n}\n\nstd::vector dcn_v2_cuda_backward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const at::Tensor &grad_output,\n int kernel_h, int kernel_w,\n int stride_h, int stride_w,\n int pad_h, int pad_w,\n int dilation_h, int dilation_w,\n int deformable_group)\n{\n\n THArgCheck(input.is_contiguous(), 1, \"input tensor has to be contiguous\");\n THArgCheck(weight.is_contiguous(), 2, \"weight tensor has to be contiguous\");\n\n AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(weight.type().is_cuda(), \"weight must be a CUDA tensor\");\n AT_ASSERTM(bias.type().is_cuda(), \"bias must be a CUDA tensor\");\n AT_ASSERTM(offset.type().is_cuda(), \"offset must be a CUDA tensor\");\n AT_ASSERTM(mask.type().is_cuda(), \"mask must be a CUDA tensor\");\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n\n const int channels_out = weight.size(0);\n const int channels_kernel = weight.size(1);\n const int kernel_h_ = weight.size(2);\n const int kernel_w_ = weight.size(3);\n\n AT_ASSERTM(kernel_h_ == kernel_h && kernel_w_ == kernel_w,\n \"Input shape and kernel shape wont match: (%d x %d vs %d x %d).\", kernel_h_, kernel_w, kernel_h_, kernel_w_);\n\n AT_ASSERTM(channels == channels_kernel,\n \"Input shape and kernel channels wont match: (%d vs %d).\", channels, channels_kernel);\n\n const int height_out = (height + 2 * pad_h - (dilation_h * (kernel_h - 1) + 1)) / stride_h + 1;\n const int width_out = (width + 2 * pad_w - (dilation_w * (kernel_w - 1) + 1)) / stride_w + 1;\n\n auto ones = at::ones({height_out, width_out}, input.options());\n auto columns = at::empty({channels * kernel_h * kernel_w, 1 * height_out * width_out}, input.options());\n auto output = at::empty({batch, channels_out, height_out, width_out}, input.options());\n\n auto grad_input = at::zeros_like(input);\n auto grad_weight = at::zeros_like(weight);\n auto grad_bias = at::zeros_like(bias);\n auto grad_offset = at::zeros_like(offset);\n auto grad_mask = at::zeros_like(mask);\n\n using scalar_t = float;\n\n for (int b = 0; b < batch; b++)\n {\n auto input_n = input.select(0, b);\n auto offset_n = offset.select(0, b);\n auto mask_n = mask.select(0, b);\n auto grad_output_n = grad_output.select(0, b);\n auto grad_input_n = grad_input.select(0, b);\n auto grad_offset_n = grad_offset.select(0, b);\n auto grad_mask_n = grad_mask.select(0, b);\n\n long m = channels * kernel_h * kernel_w;\n long n = height_out * width_out;\n long k = channels_out;\n\n THCudaBlas_Sgemm(state, 'n', 't', n, m, k, 1.0f,\n grad_output_n.data_ptr(), n,\n weight.data_ptr(), m, 0.0f,\n columns.data_ptr(), n);\n\n // gradient w.r.t. input coordinate data\n modulated_deformable_col2im_coord_cuda(c10::cuda::getCurrentCUDAStream(),\n columns.data_ptr(),\n input_n.data_ptr(),\n offset_n.data_ptr(),\n mask_n.data_ptr(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, deformable_group,\n grad_offset_n.data_ptr(),\n grad_mask_n.data_ptr());\n // gradient w.r.t. input data\n modulated_deformable_col2im_cuda(c10::cuda::getCurrentCUDAStream(),\n columns.data_ptr(),\n offset_n.data_ptr(),\n mask_n.data_ptr(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, deformable_group,\n grad_input_n.data_ptr());\n\n // gradient w.r.t. weight, dWeight should accumulate across the batch and group\n modulated_deformable_im2col_cuda(c10::cuda::getCurrentCUDAStream(),\n input_n.data_ptr(),\n offset_n.data_ptr(),\n mask_n.data_ptr(),\n 1, channels, height, width,\n height_out, width_out, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, deformable_group,\n columns.data_ptr());\n\n long m_ = channels_out;\n long n_ = channels * kernel_h * kernel_w;\n long k_ = height_out * width_out;\n\n THCudaBlas_Sgemm(state, 't', 'n', n_, m_, k_, 1.0f,\n columns.data_ptr(), k_,\n grad_output_n.data_ptr(), k_, 1.0f,\n grad_weight.data_ptr(), n_);\n\n cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();\n cublasOperation_t op = _cublasOpFromChar('t');\n _cublasAdjustLdLevel2(k_, m_, &k_);\n scalar_t* grad_output_n_float = grad_output_n.data_ptr();\n scalar_t* one_float = ones.data_ptr();\n scalar_t alpha = 1.0;\n scalar_t beta = 1.0;\n cublasSgemv(handle, op, k_, m_, &alpha, grad_output_n_float,k_, one_float,1, &beta, grad_bias.data_ptr(), 1);\n\n }\n \n\n return {\n grad_input, grad_offset, grad_mask, grad_weight, grad_bias\n };\n}\n"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cuda/dcn_v2_im2col_cuda.cu",
"content": "#include \"dcn_v2_im2col_cuda.h\"\n#include \n#include \n#include \n\n#include \n#include \n\n#include \n#include \n#include \n\n#define CUDA_KERNEL_LOOP(i, n) \\\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; \\\n i < (n); \\\n i += blockDim.x * gridDim.x)\n\nconst int CUDA_NUM_THREADS = 1024;\ninline int GET_BLOCKS(const int N)\n{\n return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;\n}\n\n\n__device__ float dmcn_im2col_bilinear_cuda(const float *bottom_data, const int data_width,\n const int height, const int width, float h, float w)\n{\n int h_low = floor(h);\n int w_low = floor(w);\n int h_high = h_low + 1;\n int w_high = w_low + 1;\n\n float lh = h - h_low;\n float lw = w - w_low;\n float hh = 1 - lh, hw = 1 - lw;\n\n float v1 = 0;\n if (h_low >= 0 && w_low >= 0)\n v1 = bottom_data[h_low * data_width + w_low];\n float v2 = 0;\n if (h_low >= 0 && w_high <= width - 1)\n v2 = bottom_data[h_low * data_width + w_high];\n float v3 = 0;\n if (h_high <= height - 1 && w_low >= 0)\n v3 = bottom_data[h_high * data_width + w_low];\n float v4 = 0;\n if (h_high <= height - 1 && w_high <= width - 1)\n v4 = bottom_data[h_high * data_width + w_high];\n\n float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;\n\n float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);\n return val;\n}\n\n__device__ float dmcn_get_gradient_weight_cuda(float argmax_h, float argmax_w,\n const int h, const int w, const int height, const int width)\n{\n if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)\n {\n //empty\n return 0;\n }\n\n int argmax_h_low = floor(argmax_h);\n int argmax_w_low = floor(argmax_w);\n int argmax_h_high = argmax_h_low + 1;\n int argmax_w_high = argmax_w_low + 1;\n\n float weight = 0;\n if (h == argmax_h_low && w == argmax_w_low)\n weight = (h + 1 - argmax_h) * (w + 1 - argmax_w);\n if (h == argmax_h_low && w == argmax_w_high)\n weight = (h + 1 - argmax_h) * (argmax_w + 1 - w);\n if (h == argmax_h_high && w == argmax_w_low)\n weight = (argmax_h + 1 - h) * (w + 1 - argmax_w);\n if (h == argmax_h_high && w == argmax_w_high)\n weight = (argmax_h + 1 - h) * (argmax_w + 1 - w);\n return weight;\n}\n\n__device__ float dmcn_get_coordinate_weight_cuda(float argmax_h, float argmax_w,\n const int height, const int width, const float *im_data,\n const int data_width, const int bp_dir)\n{\n if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width)\n {\n //empty\n return 0;\n }\n\n int argmax_h_low = floor(argmax_h);\n int argmax_w_low = floor(argmax_w);\n int argmax_h_high = argmax_h_low + 1;\n int argmax_w_high = argmax_w_low + 1;\n\n float weight = 0;\n\n if (bp_dir == 0)\n {\n if (argmax_h_low >= 0 && argmax_w_low >= 0)\n weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low];\n if (argmax_h_low >= 0 && argmax_w_high <= width - 1)\n weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high];\n if (argmax_h_high <= height - 1 && argmax_w_low >= 0)\n weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low];\n if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)\n weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high];\n }\n else if (bp_dir == 1)\n {\n if (argmax_h_low >= 0 && argmax_w_low >= 0)\n weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low];\n if (argmax_h_low >= 0 && argmax_w_high <= width - 1)\n weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high];\n if (argmax_h_high <= height - 1 && argmax_w_low >= 0)\n weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low];\n if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1)\n weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high];\n }\n\n return weight;\n}\n\n__global__ void modulated_deformable_im2col_gpu_kernel(const int n,\n const float *data_im, const float *data_offset, const float *data_mask,\n const int height, const int width, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w,\n const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int channel_per_deformable_group,\n const int batch_size, const int num_channels, const int deformable_group,\n const int height_col, const int width_col,\n float *data_col)\n{\n // launch channels * batch_size * height_col * width_col cores\n CUDA_KERNEL_LOOP(index, n)\n {\n // NOTE(CharlesShang): different from Dai Jifeng's MXNet implementation, col_buffer is of shape (c*kw*kh, N, oh, ow)\n // here columns is of shape (N, c*kw*kh, oh * ow), need to adapt axis\n\n // index index of output matrix\n const int w_col = index % width_col;\n const int h_col = (index / width_col) % height_col;\n // const int b_col = (index / width_col / height_col) % batch_size;\n const int b_col = (index / width_col / height_col / num_channels) % batch_size;\n // const int c_im = (index / width_col / height_col) / batch_size;\n const int c_im = (index / width_col / height_col) % num_channels;\n // const int c_col = c_im * kernel_h * kernel_w;\n const int c_col = c_im * kernel_h * kernel_w;\n\n // compute deformable group index\n const int deformable_group_index = c_im / channel_per_deformable_group;\n\n const int h_in = h_col * stride_h - pad_h;\n const int w_in = w_col * stride_w - pad_w;\n\n // float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col;\n float *data_col_ptr = data_col + ((b_col * num_channels * kernel_w * kernel_h + c_col) * height_col + h_col) * width_col + w_col;\n //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in;\n const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width;\n const float *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;\n\n const float *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;\n\n for (int i = 0; i < kernel_h; ++i)\n {\n for (int j = 0; j < kernel_w; ++j)\n {\n const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col;\n const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col;\n const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col;\n const float offset_h = data_offset_ptr[data_offset_h_ptr];\n const float offset_w = data_offset_ptr[data_offset_w_ptr];\n const float mask = data_mask_ptr[data_mask_hw_ptr];\n float val = static_cast(0);\n const float h_im = h_in + i * dilation_h + offset_h;\n const float w_im = w_in + j * dilation_w + offset_w;\n //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) {\n if (h_im > -1 && w_im > -1 && h_im < height && w_im < width)\n {\n //const float map_h = i * dilation_h + offset_h;\n //const float map_w = j * dilation_w + offset_w;\n //const int cur_height = height - h_in;\n //const int cur_width = width - w_in;\n //val = dmcn_im2col_bilinear_cuda(data_im_ptr, width, cur_height, cur_width, map_h, map_w);\n val = dmcn_im2col_bilinear_cuda(data_im_ptr, width, height, width, h_im, w_im);\n }\n *data_col_ptr = val * mask;\n // data_col_ptr += batch_size * height_col * width_col;\n data_col_ptr += height_col * width_col;\n }\n }\n }\n}\n\n__global__ void modulated_deformable_col2im_gpu_kernel(const int n,\n const float *data_col, const float *data_offset, const float *data_mask,\n const int channels, const int height, const int width,\n const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w,\n const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int channel_per_deformable_group,\n const int batch_size, const int deformable_group,\n const int height_col, const int width_col,\n float *grad_im)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n const int j = (index / width_col / height_col / batch_size) % kernel_w;\n const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h;\n const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h;\n // compute the start and end of the output\n\n const int deformable_group_index = c / channel_per_deformable_group;\n\n int w_out = index % width_col;\n int h_out = (index / width_col) % height_col;\n int b = (index / width_col / height_col) % batch_size;\n int w_in = w_out * stride_w - pad_w;\n int h_in = h_out * stride_h - pad_h;\n\n const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;\n const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;\n const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out;\n const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out;\n const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_out) * width_col + w_out;\n const float offset_h = data_offset_ptr[data_offset_h_ptr];\n const float offset_w = data_offset_ptr[data_offset_w_ptr];\n const float mask = data_mask_ptr[data_mask_hw_ptr];\n const float cur_inv_h_data = h_in + i * dilation_h + offset_h;\n const float cur_inv_w_data = w_in + j * dilation_w + offset_w;\n\n const float cur_top_grad = data_col[index] * mask;\n const int cur_h = (int)cur_inv_h_data;\n const int cur_w = (int)cur_inv_w_data;\n for (int dy = -2; dy <= 2; dy++)\n {\n for (int dx = -2; dx <= 2; dx++)\n {\n if (cur_h + dy >= 0 && cur_h + dy < height &&\n cur_w + dx >= 0 && cur_w + dx < width &&\n abs(cur_inv_h_data - (cur_h + dy)) < 1 &&\n abs(cur_inv_w_data - (cur_w + dx)) < 1)\n {\n int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx;\n float weight = dmcn_get_gradient_weight_cuda(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width);\n atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad);\n }\n }\n }\n }\n}\n\n__global__ void modulated_deformable_col2im_coord_gpu_kernel(const int n,\n const float *data_col, const float *data_im,\n const float *data_offset, const float *data_mask,\n const int channels, const int height, const int width,\n const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w,\n const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int channel_per_deformable_group,\n const int batch_size, const int offset_channels, const int deformable_group,\n const int height_col, const int width_col,\n float *grad_offset, float *grad_mask)\n{\n CUDA_KERNEL_LOOP(index, n)\n {\n float val = 0, mval = 0;\n int w = index % width_col;\n int h = (index / width_col) % height_col;\n int c = (index / width_col / height_col) % offset_channels;\n int b = (index / width_col / height_col) / offset_channels;\n // compute the start and end of the output\n\n const int deformable_group_index = c / (2 * kernel_h * kernel_w);\n const int col_step = kernel_h * kernel_w;\n int cnt = 0;\n const float *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group * batch_size * width_col * height_col;\n const float *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) * channel_per_deformable_group / kernel_h / kernel_w * height * width;\n const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col;\n const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col;\n\n const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w;\n\n for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step)\n {\n const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w;\n const int bp_dir = offset_c % 2;\n\n int j = (col_pos / width_col / height_col / batch_size) % kernel_w;\n int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h;\n int w_out = col_pos % width_col;\n int h_out = (col_pos / width_col) % height_col;\n int w_in = w_out * stride_w - pad_w;\n int h_in = h_out * stride_h - pad_h;\n const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out);\n const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out);\n const int data_mask_hw_ptr = (((i * kernel_w + j) * height_col + h_out) * width_col + w_out);\n const float offset_h = data_offset_ptr[data_offset_h_ptr];\n const float offset_w = data_offset_ptr[data_offset_w_ptr];\n const float mask = data_mask_ptr[data_mask_hw_ptr];\n float inv_h = h_in + i * dilation_h + offset_h;\n float inv_w = w_in + j * dilation_w + offset_w;\n if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width)\n {\n inv_h = inv_w = -2;\n }\n else\n {\n mval += data_col_ptr[col_pos] * dmcn_im2col_bilinear_cuda(data_im_ptr + cnt * height * width, width, height, width, inv_h, inv_w);\n }\n const float weight = dmcn_get_coordinate_weight_cuda(\n inv_h, inv_w,\n height, width, data_im_ptr + cnt * height * width, width, bp_dir);\n val += weight * data_col_ptr[col_pos] * mask;\n cnt += 1;\n }\n // KERNEL_ASSIGN(grad_offset[index], offset_req, val);\n grad_offset[index] = val;\n if (offset_c % 2 == 0)\n // KERNEL_ASSIGN(grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w], mask_req, mval);\n grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w] = mval;\n }\n}\n\nvoid modulated_deformable_im2col_cuda(cudaStream_t stream,\n const float* data_im, const float* data_offset, const float* data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im, \n const int height_col, const int width_col, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w, \n const int dilation_h, const int dilation_w,\n const int deformable_group, float* data_col) {\n // num_axes should be smaller than block size\n const int channel_per_deformable_group = channels / deformable_group;\n const int num_kernels = channels * batch_size * height_col * width_col;\n modulated_deformable_im2col_gpu_kernel\n <<>>(\n num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kernel_w,\n pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group,\n batch_size, channels, deformable_group, height_col, width_col, data_col);\n \n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in modulated_deformable_im2col_cuda: %s\\n\", cudaGetErrorString(err));\n }\n\n}\n\nvoid modulated_deformable_col2im_cuda(cudaStream_t stream,\n const float* data_col, const float* data_offset, const float* data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im, \n const int height_col, const int width_col, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w, \n const int dilation_h, const int dilation_w, \n const int deformable_group, float* grad_im){\n\n const int channel_per_deformable_group = channels / deformable_group;\n const int num_kernels = channels * kernel_h * kernel_w * batch_size * height_col * width_col;\n modulated_deformable_col2im_gpu_kernel\n <<>>(\n num_kernels, data_col, data_offset, data_mask, channels, height_im, width_im,\n kernel_h, kernel_w, pad_h, pad_h, stride_h, stride_w,\n dilation_h, dilation_w, channel_per_deformable_group,\n batch_size, deformable_group, height_col, width_col, grad_im);\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in modulated_deformable_col2im_cuda: %s\\n\", cudaGetErrorString(err));\n }\n\n}\n\nvoid modulated_deformable_col2im_coord_cuda(cudaStream_t stream,\n const float* data_col, const float* data_im, const float* data_offset, const float* data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im, \n const int height_col, const int width_col, const int kernel_h, const int kernel_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w, \n const int dilation_h, const int dilation_w, \n const int deformable_group,\n float* grad_offset, float* grad_mask) {\n const int num_kernels = batch_size * height_col * width_col * 2 * kernel_h * kernel_w * deformable_group;\n const int channel_per_deformable_group = channels * kernel_h * kernel_w / deformable_group;\n modulated_deformable_col2im_coord_gpu_kernel\n <<>>(\n num_kernels, data_col, data_im, data_offset, data_mask, channels, height_im, width_im,\n kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w,\n dilation_h, dilation_w, channel_per_deformable_group,\n batch_size, 2 * kernel_h * kernel_w * deformable_group, deformable_group, height_col, width_col, \n grad_offset, grad_mask);\n cudaError_t err = cudaGetLastError();\n if (err != cudaSuccess)\n {\n printf(\"error in modulated_deformable_col2im_coord_cuda: %s\\n\", cudaGetErrorString(err));\n }\n}"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cuda/dcn_v2_im2col_cuda.h",
"content": "\n/*!\n ******************* BEGIN Caffe Copyright Notice and Disclaimer ****************\n *\n * COPYRIGHT\n *\n * All contributions by the University of California:\n * Copyright (c) 2014-2017 The Regents of the University of California (Regents)\n * All rights reserved.\n *\n * All other contributions:\n * Copyright (c) 2014-2017, the respective contributors\n * All rights reserved.\n *\n * Caffe uses a shared copyright model: each contributor holds copyright over\n * their contributions to Caffe. The project versioning records all such\n * contribution and copyright details. If a contributor wants to further mark\n * their specific copyright on a particular contribution, they should indicate\n * their copyright solely in the commit message of the change when it is\n * committed.\n *\n * LICENSE\n *\n * Redistribution and use in source and binary forms, with or without\n * modification, are permitted provided that the following conditions are met:\n *\n * 1. Redistributions of source code must retain the above copyright notice, this\n * list of conditions and the following disclaimer.\n * 2. Redistributions in binary form must reproduce the above copyright notice,\n * this list of conditions and the following disclaimer in the documentation\n * and/or other materials provided with the distribution.\n *\n * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\" AND\n * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED\n * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE\n * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR\n * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES\n * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;\n * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND\n * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT\n * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS\n * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.\n *\n * CONTRIBUTION AGREEMENT\n *\n * By contributing to the BVLC/caffe repository through pull-request, comment,\n * or otherwise, the contributor releases their content to the\n * license and copyright terms herein.\n *\n ***************** END Caffe Copyright Notice and Disclaimer ********************\n *\n * Copyright (c) 2018 Microsoft\n * Licensed under The MIT License [see LICENSE for details]\n * \\file modulated_deformable_im2col.h\n * \\brief Function definitions of converting an image to\n * column matrix based on kernel, padding, dilation, and offset.\n * These functions are mainly used in deformable convolution operators.\n * \\ref: https://arxiv.org/abs/1811.11168\n * \\author Yuwen Xiong, Haozhi Qi, Jifeng Dai, Xizhou Zhu, Han Hu\n */\n\n/***************** Adapted by Charles Shang *********************/\n\n#ifndef DCN_V2_IM2COL_CUDA\n#define DCN_V2_IM2COL_CUDA\n\n#ifdef __cplusplus\nextern \"C\"\n{\n#endif\n\n void modulated_deformable_im2col_cuda(cudaStream_t stream,\n const float *data_im, const float *data_offset, const float *data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im,\n const int height_col, const int width_col, const int kernel_h, const int kenerl_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int deformable_group, float *data_col);\n\n void modulated_deformable_col2im_cuda(cudaStream_t stream,\n const float *data_col, const float *data_offset, const float *data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im,\n const int height_col, const int width_col, const int kernel_h, const int kenerl_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int deformable_group, float *grad_im);\n\n void modulated_deformable_col2im_coord_cuda(cudaStream_t stream,\n const float *data_col, const float *data_im, const float *data_offset, const float *data_mask,\n const int batch_size, const int channels, const int height_im, const int width_im,\n const int height_col, const int width_col, const int kernel_h, const int kenerl_w,\n const int pad_h, const int pad_w, const int stride_h, const int stride_w,\n const int dilation_h, const int dilation_w,\n const int deformable_group,\n float *grad_offset, float *grad_mask);\n\n#ifdef __cplusplus\n}\n#endif\n\n#endif"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cuda/dcn_v2_psroi_pooling_cuda.cu",
"content": "/*!\n * Copyright (c) 2017 Microsoft\n * Licensed under The MIT License [see LICENSE for details]\n * \\file deformable_psroi_pooling.cu\n * \\brief\n * \\author Yi Li, Guodong Zhang, Jifeng Dai\n*/\n/***************** Adapted by Charles Shang *********************/\n\n#include \n#include \n#include \n#include \n\n#include \n#include \n\n#include \n#include \n#include \n\n#define CUDA_KERNEL_LOOP(i, n) \\\n for (int i = blockIdx.x * blockDim.x + threadIdx.x; \\\n i < (n); \\\n i += blockDim.x * gridDim.x)\n\nconst int CUDA_NUM_THREADS = 1024;\ninline int GET_BLOCKS(const int N)\n{\n return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS;\n}\n\ntemplate \n__device__ T bilinear_interp_cuda(\n const T *data,\n const T x,\n const T y,\n const int width,\n const int height)\n{\n int x1 = floor(x);\n int x2 = ceil(x);\n int y1 = floor(y);\n int y2 = ceil(y);\n T dist_x = static_cast(x - x1);\n T dist_y = static_cast(y - y1);\n T value11 = data[y1 * width + x1];\n T value12 = data[y2 * width + x1];\n T value21 = data[y1 * width + x2];\n T value22 = data[y2 * width + x2];\n T value = (1 - dist_x) * (1 - dist_y) * value11 +\n (1 - dist_x) * dist_y * value12 +\n dist_x * (1 - dist_y) * value21 +\n dist_x * dist_y * value22;\n return value;\n}\n\ntemplate \n__global__ void DeformablePSROIPoolForwardKernelCuda(\n const int count,\n const T *bottom_data,\n const T spatial_scale,\n const int channels,\n const int height, const int width,\n const int pooled_height, const int pooled_width,\n const T *bottom_rois, const T *bottom_trans,\n const int no_trans,\n const T trans_std,\n const int sample_per_part,\n const int output_dim,\n const int group_size,\n const int part_size,\n const int num_classes,\n const int channels_each_class,\n T *top_data,\n T *top_count)\n{\n CUDA_KERNEL_LOOP(index, count)\n {\n // The output is in order (n, ctop, ph, pw)\n int pw = index % pooled_width;\n int ph = (index / pooled_width) % pooled_height;\n int ctop = (index / pooled_width / pooled_height) % output_dim;\n int n = index / pooled_width / pooled_height / output_dim;\n\n // [start, end) interval for spatial sampling\n const T *offset_bottom_rois = bottom_rois + n * 5;\n int roi_batch_ind = offset_bottom_rois[0];\n T roi_start_w = static_cast(round(offset_bottom_rois[1])) * spatial_scale - 0.5;\n T roi_start_h = static_cast(round(offset_bottom_rois[2])) * spatial_scale - 0.5;\n T roi_end_w = static_cast(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;\n T roi_end_h = static_cast(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;\n\n // Force too small ROIs to be 1x1\n T roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0\n T roi_height = max(roi_end_h - roi_start_h, 0.1);\n\n // Compute w and h at bottom\n T bin_size_h = roi_height / static_cast(pooled_height);\n T bin_size_w = roi_width / static_cast(pooled_width);\n\n T sub_bin_size_h = bin_size_h / static_cast(sample_per_part);\n T sub_bin_size_w = bin_size_w / static_cast(sample_per_part);\n\n int part_h = floor(static_cast(ph) / pooled_height * part_size);\n int part_w = floor(static_cast(pw) / pooled_width * part_size);\n int class_id = ctop / channels_each_class;\n T trans_x = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * trans_std;\n T trans_y = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * trans_std;\n\n T wstart = static_cast(pw) * bin_size_w + roi_start_w;\n wstart += trans_x * roi_width;\n T hstart = static_cast(ph) * bin_size_h + roi_start_h;\n hstart += trans_y * roi_height;\n\n T sum = 0;\n int count = 0;\n int gw = floor(static_cast(pw) * group_size / pooled_width);\n int gh = floor(static_cast(ph) * group_size / pooled_height);\n gw = min(max(gw, 0), group_size - 1);\n gh = min(max(gh, 0), group_size - 1);\n\n const T *offset_bottom_data = bottom_data + (roi_batch_ind * channels) * height * width;\n for (int ih = 0; ih < sample_per_part; ih++)\n {\n for (int iw = 0; iw < sample_per_part; iw++)\n {\n T w = wstart + iw * sub_bin_size_w;\n T h = hstart + ih * sub_bin_size_h;\n // bilinear interpolation\n if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)\n {\n continue;\n }\n w = min(max(w, 0.), width - 1.);\n h = min(max(h, 0.), height - 1.);\n int c = (ctop * group_size + gh) * group_size + gw;\n T val = bilinear_interp_cuda(offset_bottom_data + c * height * width, w, h, width, height);\n sum += val;\n count++;\n }\n }\n top_data[index] = count == 0 ? static_cast(0) : sum / count;\n top_count[index] = count;\n }\n}\n\ntemplate \n__global__ void DeformablePSROIPoolBackwardAccKernelCuda(\n const int count,\n const T *top_diff,\n const T *top_count,\n const int num_rois,\n const T spatial_scale,\n const int channels,\n const int height, const int width,\n const int pooled_height, const int pooled_width,\n const int output_dim,\n T *bottom_data_diff, T *bottom_trans_diff,\n const T *bottom_data,\n const T *bottom_rois,\n const T *bottom_trans,\n const int no_trans,\n const T trans_std,\n const int sample_per_part,\n const int group_size,\n const int part_size,\n const int num_classes,\n const int channels_each_class)\n{\n CUDA_KERNEL_LOOP(index, count)\n {\n // The output is in order (n, ctop, ph, pw)\n int pw = index % pooled_width;\n int ph = (index / pooled_width) % pooled_height;\n int ctop = (index / pooled_width / pooled_height) % output_dim;\n int n = index / pooled_width / pooled_height / output_dim;\n\n // [start, end) interval for spatial sampling\n const T *offset_bottom_rois = bottom_rois + n * 5;\n int roi_batch_ind = offset_bottom_rois[0];\n T roi_start_w = static_cast(round(offset_bottom_rois[1])) * spatial_scale - 0.5;\n T roi_start_h = static_cast(round(offset_bottom_rois[2])) * spatial_scale - 0.5;\n T roi_end_w = static_cast(round(offset_bottom_rois[3]) + 1.) * spatial_scale - 0.5;\n T roi_end_h = static_cast(round(offset_bottom_rois[4]) + 1.) * spatial_scale - 0.5;\n\n // Force too small ROIs to be 1x1\n T roi_width = max(roi_end_w - roi_start_w, 0.1); //avoid 0\n T roi_height = max(roi_end_h - roi_start_h, 0.1);\n\n // Compute w and h at bottom\n T bin_size_h = roi_height / static_cast(pooled_height);\n T bin_size_w = roi_width / static_cast(pooled_width);\n\n T sub_bin_size_h = bin_size_h / static_cast(sample_per_part);\n T sub_bin_size_w = bin_size_w / static_cast(sample_per_part);\n\n int part_h = floor(static_cast(ph) / pooled_height * part_size);\n int part_w = floor(static_cast(pw) / pooled_width * part_size);\n int class_id = ctop / channels_each_class;\n T trans_x = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w] * trans_std;\n T trans_y = no_trans ? static_cast(0) : bottom_trans[(((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w] * trans_std;\n\n T wstart = static_cast(pw) * bin_size_w + roi_start_w;\n wstart += trans_x * roi_width;\n T hstart = static_cast(ph) * bin_size_h + roi_start_h;\n hstart += trans_y * roi_height;\n\n if (top_count[index] <= 0)\n {\n continue;\n }\n T diff_val = top_diff[index] / top_count[index];\n const T *offset_bottom_data = bottom_data + roi_batch_ind * channels * height * width;\n T *offset_bottom_data_diff = bottom_data_diff + roi_batch_ind * channels * height * width;\n int gw = floor(static_cast(pw) * group_size / pooled_width);\n int gh = floor(static_cast(ph) * group_size / pooled_height);\n gw = min(max(gw, 0), group_size - 1);\n gh = min(max(gh, 0), group_size - 1);\n\n for (int ih = 0; ih < sample_per_part; ih++)\n {\n for (int iw = 0; iw < sample_per_part; iw++)\n {\n T w = wstart + iw * sub_bin_size_w;\n T h = hstart + ih * sub_bin_size_h;\n // bilinear interpolation\n if (w < -0.5 || w > width - 0.5 || h < -0.5 || h > height - 0.5)\n {\n continue;\n }\n w = min(max(w, 0.), width - 1.);\n h = min(max(h, 0.), height - 1.);\n int c = (ctop * group_size + gh) * group_size + gw;\n // backward on feature\n int x0 = floor(w);\n int x1 = ceil(w);\n int y0 = floor(h);\n int y1 = ceil(h);\n T dist_x = w - x0, dist_y = h - y0;\n T q00 = (1 - dist_x) * (1 - dist_y);\n T q01 = (1 - dist_x) * dist_y;\n T q10 = dist_x * (1 - dist_y);\n T q11 = dist_x * dist_y;\n int bottom_index_base = c * height * width;\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x0, q00 * diff_val);\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x0, q01 * diff_val);\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y0 * width + x1, q10 * diff_val);\n atomicAdd(offset_bottom_data_diff + bottom_index_base + y1 * width + x1, q11 * diff_val);\n\n if (no_trans)\n {\n continue;\n }\n T U00 = offset_bottom_data[bottom_index_base + y0 * width + x0];\n T U01 = offset_bottom_data[bottom_index_base + y1 * width + x0];\n T U10 = offset_bottom_data[bottom_index_base + y0 * width + x1];\n T U11 = offset_bottom_data[bottom_index_base + y1 * width + x1];\n T diff_x = (U11 * dist_y + U10 * (1 - dist_y) - U01 * dist_y - U00 * (1 - dist_y)) * trans_std * diff_val;\n diff_x *= roi_width;\n T diff_y = (U11 * dist_x + U01 * (1 - dist_x) - U10 * dist_x - U00 * (1 - dist_x)) * trans_std * diff_val;\n diff_y *= roi_height;\n\n atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2) * part_size + part_h) * part_size + part_w, diff_x);\n atomicAdd(bottom_trans_diff + (((n * num_classes + class_id) * 2 + 1) * part_size + part_h) * part_size + part_w, diff_y);\n }\n }\n }\n}\n\nstd::tuple\ndcn_v2_psroi_pooling_cuda_forward(const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std)\n{\n AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(bbox.type().is_cuda(), \"rois must be a CUDA tensor\");\n AT_ASSERTM(trans.type().is_cuda(), \"trans must be a CUDA tensor\");\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n const int channels_trans = no_trans ? 2 : trans.size(1);\n const int num_bbox = bbox.size(0);\n\n AT_ASSERTM(channels == output_dim, \"input channels and output channels must equal\");\n auto pooled_height = pooled_size;\n auto pooled_width = pooled_size;\n\n auto out = at::empty({num_bbox, output_dim, pooled_height, pooled_width}, input.options());\n long out_size = num_bbox * output_dim * pooled_height * pooled_width;\n auto top_count = at::zeros({num_bbox, output_dim, pooled_height, pooled_width}, input.options());\n\n const int num_classes = no_trans ? 1 : channels_trans / 2;\n const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;\n\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n if (out.numel() == 0)\n {\n THCudaCheck(cudaGetLastError());\n return std::make_tuple(out, top_count);\n }\n\n dim3 grid(std::min(THCCeilDiv(out_size, 512L), 4096L));\n dim3 block(512);\n\n AT_DISPATCH_FLOATING_TYPES(input.type(), \"dcn_v2_psroi_pooling_cuda_forward\", [&] {\n DeformablePSROIPoolForwardKernelCuda<<>>(\n out_size,\n input.contiguous().data_ptr(),\n spatial_scale,\n channels,\n height, width,\n pooled_height,\n pooled_width,\n bbox.contiguous().data_ptr(),\n trans.contiguous().data_ptr(),\n no_trans,\n trans_std,\n sample_per_part,\n output_dim,\n group_size,\n part_size,\n num_classes,\n channels_each_class,\n out.data_ptr(),\n top_count.data_ptr());\n });\n THCudaCheck(cudaGetLastError());\n return std::make_tuple(out, top_count);\n}\n\nstd::tuple\ndcn_v2_psroi_pooling_cuda_backward(const at::Tensor &out_grad,\n const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const at::Tensor &top_count,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std)\n{\n AT_ASSERTM(out_grad.type().is_cuda(), \"out_grad must be a CUDA tensor\");\n AT_ASSERTM(input.type().is_cuda(), \"input must be a CUDA tensor\");\n AT_ASSERTM(bbox.type().is_cuda(), \"bbox must be a CUDA tensor\");\n AT_ASSERTM(trans.type().is_cuda(), \"trans must be a CUDA tensor\");\n AT_ASSERTM(top_count.type().is_cuda(), \"top_count must be a CUDA tensor\");\n\n const int batch = input.size(0);\n const int channels = input.size(1);\n const int height = input.size(2);\n const int width = input.size(3);\n const int channels_trans = no_trans ? 2 : trans.size(1);\n const int num_bbox = bbox.size(0);\n\n AT_ASSERTM(channels == output_dim, \"input channels and output channels must equal\");\n auto pooled_height = pooled_size;\n auto pooled_width = pooled_size;\n long out_size = num_bbox * output_dim * pooled_height * pooled_width;\n const int num_classes = no_trans ? 1 : channels_trans / 2;\n const int channels_each_class = no_trans ? output_dim : output_dim / num_classes;\n\n auto input_grad = at::zeros({batch, channels, height, width}, out_grad.options());\n auto trans_grad = at::zeros_like(trans);\n\n if (input_grad.numel() == 0)\n {\n THCudaCheck(cudaGetLastError());\n return std::make_tuple(input_grad, trans_grad);\n }\n\n dim3 grid(std::min(THCCeilDiv(out_size, 512L), 4096L));\n dim3 block(512);\n cudaStream_t stream = at::cuda::getCurrentCUDAStream();\n\n AT_DISPATCH_FLOATING_TYPES(out_grad.type(), \"dcn_v2_psroi_pooling_cuda_backward\", [&] {\n DeformablePSROIPoolBackwardAccKernelCuda<<>>(\n out_size,\n out_grad.contiguous().data_ptr(),\n top_count.contiguous().data_ptr(),\n num_bbox,\n spatial_scale,\n channels,\n height,\n width,\n pooled_height,\n pooled_width,\n output_dim,\n input_grad.contiguous().data_ptr(),\n trans_grad.contiguous().data_ptr(),\n input.contiguous().data_ptr(),\n bbox.contiguous().data_ptr(),\n trans.contiguous().data_ptr(),\n no_trans,\n trans_std,\n sample_per_part,\n group_size,\n part_size,\n num_classes,\n channels_each_class);\n });\n THCudaCheck(cudaGetLastError());\n return std::make_tuple(input_grad, trans_grad);\n}"
},
{
"path": "code/real/bsrt/model/DCNv2/src/cuda/vision.h",
"content": "#pragma once\n#include \n#include \nat::Tensor\ndcn_v2_cuda_forward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const int kernel_h,\n const int kernel_w,\n const int stride_h,\n const int stride_w,\n const int pad_h,\n const int pad_w,\n const int dilation_h,\n const int dilation_w,\n const int deformable_group);\n\nstd::vector\ndcn_v2_cuda_backward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const at::Tensor &grad_output,\n int kernel_h, int kernel_w,\n int stride_h, int stride_w,\n int pad_h, int pad_w,\n int dilation_h, int dilation_w,\n int deformable_group);\n\n\nstd::tuple\ndcn_v2_psroi_pooling_cuda_forward(const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std);\n\nstd::tuple\ndcn_v2_psroi_pooling_cuda_backward(const at::Tensor &out_grad,\n const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const at::Tensor &top_count,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std);"
},
{
"path": "code/real/bsrt/model/DCNv2/src/dcn_v2.h",
"content": "#pragma once\n\n#include \"cpu/vision.h\"\n\n#ifdef WITH_CUDA\n#include \"cuda/vision.h\"\n#endif\n\nat::Tensor\ndcn_v2_forward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const int kernel_h,\n const int kernel_w,\n const int stride_h,\n const int stride_w,\n const int pad_h,\n const int pad_w,\n const int dilation_h,\n const int dilation_w,\n const int deformable_group)\n{\n if (input.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return dcn_v2_cuda_forward(input, weight, bias, offset, mask,\n kernel_h, kernel_w,\n stride_h, stride_w,\n pad_h, pad_w,\n dilation_h, dilation_w,\n deformable_group);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n else{\n return dcn_v2_cpu_forward(input, weight, bias, offset, mask,\n kernel_h, kernel_w,\n stride_h, stride_w,\n pad_h, pad_w,\n dilation_h, dilation_w,\n deformable_group);\n }\n}\n\nstd::vector\ndcn_v2_backward(const at::Tensor &input,\n const at::Tensor &weight,\n const at::Tensor &bias,\n const at::Tensor &offset,\n const at::Tensor &mask,\n const at::Tensor &grad_output,\n int kernel_h, int kernel_w,\n int stride_h, int stride_w,\n int pad_h, int pad_w,\n int dilation_h, int dilation_w,\n int deformable_group)\n{\n if (input.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return dcn_v2_cuda_backward(input,\n weight,\n bias,\n offset,\n mask,\n grad_output,\n kernel_h, kernel_w,\n stride_h, stride_w,\n pad_h, pad_w,\n dilation_h, dilation_w,\n deformable_group);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n else{\n return dcn_v2_cpu_backward(input,\n weight,\n bias,\n offset,\n mask,\n grad_output,\n kernel_h, kernel_w,\n stride_h, stride_w,\n pad_h, pad_w,\n dilation_h, dilation_w,\n deformable_group);\n }\n}\n\nstd::tuple\ndcn_v2_psroi_pooling_forward(const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std)\n{\n if (input.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return dcn_v2_psroi_pooling_cuda_forward(input,\n bbox,\n trans,\n no_trans,\n spatial_scale,\n output_dim,\n group_size,\n pooled_size,\n part_size,\n sample_per_part,\n trans_std);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n else{\n return dcn_v2_psroi_pooling_cpu_forward(input,\n bbox,\n trans,\n no_trans,\n spatial_scale,\n output_dim,\n group_size,\n pooled_size,\n part_size,\n sample_per_part,\n trans_std);\n }\n}\n\nstd::tuple\ndcn_v2_psroi_pooling_backward(const at::Tensor &out_grad,\n const at::Tensor &input,\n const at::Tensor &bbox,\n const at::Tensor &trans,\n const at::Tensor &top_count,\n const int no_trans,\n const float spatial_scale,\n const int output_dim,\n const int group_size,\n const int pooled_size,\n const int part_size,\n const int sample_per_part,\n const float trans_std)\n{\n if (input.type().is_cuda())\n {\n#ifdef WITH_CUDA\n return dcn_v2_psroi_pooling_cuda_backward(out_grad,\n input,\n bbox,\n trans,\n top_count,\n no_trans,\n spatial_scale,\n output_dim,\n group_size,\n pooled_size,\n part_size,\n sample_per_part,\n trans_std);\n#else\n AT_ERROR(\"Not compiled with GPU support\");\n#endif\n }\n else{\n return dcn_v2_psroi_pooling_cpu_backward(out_grad,\n input,\n bbox,\n trans,\n top_count,\n no_trans,\n spatial_scale,\n output_dim,\n group_size,\n pooled_size,\n part_size,\n sample_per_part,\n trans_std);\n }\n}"
},
{
"path": "code/real/bsrt/model/DCNv2/src/vision.cpp",
"content": "\n#include \"dcn_v2.h\"\n\nPYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {\n m.def(\"dcn_v2_forward\", &dcn_v2_forward, \"dcn_v2_forward\");\n m.def(\"dcn_v2_backward\", &dcn_v2_backward, \"dcn_v2_backward\");\n m.def(\"dcn_v2_psroi_pooling_forward\", &dcn_v2_psroi_pooling_forward, \"dcn_v2_psroi_pooling_forward\");\n m.def(\"dcn_v2_psroi_pooling_backward\", &dcn_v2_psroi_pooling_backward, \"dcn_v2_psroi_pooling_backward\");\n}\n"
},
{
"path": "code/real/bsrt/model/DCNv2/test.py",
"content": "#!/usr/bin/env python\nfrom __future__ import absolute_import\nfrom __future__ import print_function\nfrom __future__ import division\n\nimport time\nimport torch\nimport torch.nn as nn\nfrom torch.autograd import gradcheck\n\nfrom dcn_v2 import dcn_v2_conv, DCNv2, DCN\nfrom dcn_v2 import dcn_v2_pooling, DCNv2Pooling, DCNPooling\n\ndeformable_groups = 1\nN, inC, inH, inW = 2, 2, 4, 4\noutC = 2\nkH, kW = 3, 3\n\n\ndef conv_identify(weight, bias):\n weight.data.zero_()\n bias.data.zero_()\n o, i, h, w = weight.shape\n y = h//2\n x = w//2\n for p in range(i):\n for q in range(o):\n if p == q:\n weight.data[q, p, y, x] = 1.0\n\n\ndef check_zero_offset():\n conv_offset = nn.Conv2d(inC, deformable_groups * 2 * kH * kW,\n kernel_size=(kH, kW),\n stride=(1, 1),\n padding=(1, 1),\n bias=True).cuda()\n\n conv_mask = nn.Conv2d(inC, deformable_groups * 1 * kH * kW,\n kernel_size=(kH, kW),\n stride=(1, 1),\n padding=(1, 1),\n bias=True).cuda()\n\n dcn_v2 = DCNv2(inC, outC, (kH, kW),\n stride=1, padding=1, dilation=1,\n deformable_groups=deformable_groups).cuda()\n\n conv_offset.weight.data.zero_()\n conv_offset.bias.data.zero_()\n conv_mask.weight.data.zero_()\n conv_mask.bias.data.zero_()\n conv_identify(dcn_v2.weight, dcn_v2.bias)\n\n input = torch.randn(N, inC, inH, inW).cuda()\n offset = conv_offset(input)\n mask = conv_mask(input)\n mask = torch.sigmoid(mask)\n output = dcn_v2(input, offset, mask)\n output *= 2\n d = (input - output).abs().max()\n if d < 1e-10:\n print('Zero offset passed')\n else:\n print('Zero offset failed')\n print(input)\n print(output)\n\ndef check_gradient_dconv():\n\n input = torch.rand(N, inC, inH, inW).cuda() * 0.01\n input.requires_grad = True\n\n offset = torch.randn(N, deformable_groups * 2 * kW * kH, inH, inW).cuda() * 2\n # offset.data.zero_()\n # offset.data -= 0.5\n offset.requires_grad = True\n\n mask = torch.rand(N, deformable_groups * 1 * kW * kH, inH, inW).cuda()\n # mask.data.zero_()\n mask.requires_grad = True\n mask = torch.sigmoid(mask)\n\n weight = torch.randn(outC, inC, kH, kW).cuda()\n weight.requires_grad = True\n\n bias = torch.rand(outC).cuda()\n bias.requires_grad = True\n\n stride = 1\n padding = 1\n dilation = 1\n\n print('check_gradient_dconv: ',\n gradcheck(dcn_v2_conv, (input, offset, mask, weight, bias,\n stride, padding, dilation, deformable_groups),\n eps=1e-3, atol=1e-4, rtol=1e-2))\n\n\ndef check_pooling_zero_offset():\n\n input = torch.randn(2, 16, 64, 64).cuda().zero_()\n input[0, :, 16:26, 16:26] = 1.\n input[1, :, 10:20, 20:30] = 2.\n rois = torch.tensor([\n [0, 65, 65, 103, 103],\n [1, 81, 41, 119, 79],\n ]).cuda().float()\n pooling = DCNv2Pooling(spatial_scale=1.0 / 4,\n pooled_size=7,\n output_dim=16,\n no_trans=True,\n group_size=1,\n trans_std=0.0).cuda()\n\n out = pooling(input, rois, input.new())\n s = ', '.join(['%f' % out[i, :, :, :].mean().item()\n for i in range(rois.shape[0])])\n print(s)\n\n dpooling = DCNv2Pooling(spatial_scale=1.0 / 4,\n pooled_size=7,\n output_dim=16,\n no_trans=False,\n group_size=1,\n trans_std=0.0).cuda()\n offset = torch.randn(20, 2, 7, 7).cuda().zero_()\n dout = dpooling(input, rois, offset)\n s = ', '.join(['%f' % dout[i, :, :, :].mean().item()\n for i in range(rois.shape[0])])\n print(s)\n\n\ndef check_gradient_dpooling():\n input = torch.randn(2, 3, 5, 5).cuda() * 0.01\n N = 4\n batch_inds = torch.randint(2, (N, 1)).cuda().float()\n x = torch.rand((N, 1)).cuda().float() * 15\n y = torch.rand((N, 1)).cuda().float() * 15\n w = torch.rand((N, 1)).cuda().float() * 10\n h = torch.rand((N, 1)).cuda().float() * 10\n rois = torch.cat((batch_inds, x, y, x + w, y + h), dim=1)\n offset = torch.randn(N, 2, 3, 3).cuda()\n input.requires_grad = True\n offset.requires_grad = True\n\n spatial_scale = 1.0 / 4\n pooled_size = 3\n output_dim = 3\n no_trans = 0\n group_size = 1\n trans_std = 0.0\n sample_per_part = 4\n part_size = pooled_size\n\n print('check_gradient_dpooling:',\n gradcheck(dcn_v2_pooling, (input, rois, offset,\n spatial_scale,\n pooled_size,\n output_dim,\n no_trans,\n group_size,\n part_size,\n sample_per_part,\n trans_std),\n eps=1e-4))\n\n\ndef example_dconv():\n input = torch.randn(2, 64, 128, 128).cuda()\n # wrap all things (offset and mask) in DCN\n dcn = DCN(64, 64, kernel_size=(3, 3), stride=1,\n padding=1, deformable_groups=2).cuda()\n # print(dcn.weight.shape, input.shape)\n output = dcn(input)\n targert = output.new(*output.size())\n targert.data.uniform_(-0.01, 0.01)\n error = (targert - output).mean()\n error.backward()\n print(output.shape)\n\n\ndef example_dpooling():\n input = torch.randn(2, 32, 64, 64).cuda()\n batch_inds = torch.randint(2, (20, 1)).cuda().float()\n x = torch.randint(256, (20, 1)).cuda().float()\n y = torch.randint(256, (20, 1)).cuda().float()\n w = torch.randint(64, (20, 1)).cuda().float()\n h = torch.randint(64, (20, 1)).cuda().float()\n rois = torch.cat((batch_inds, x, y, x + w, y + h), dim=1)\n offset = torch.randn(20, 2, 7, 7).cuda()\n input.requires_grad = True\n offset.requires_grad = True\n\n # normal roi_align\n pooling = DCNv2Pooling(spatial_scale=1.0 / 4,\n pooled_size=7,\n output_dim=32,\n no_trans=True,\n group_size=1,\n trans_std=0.1).cuda()\n\n # deformable pooling\n dpooling = DCNv2Pooling(spatial_scale=1.0 / 4,\n pooled_size=7,\n output_dim=32,\n no_trans=False,\n group_size=1,\n trans_std=0.1).cuda()\n\n out = pooling(input, rois, offset)\n dout = dpooling(input, rois, offset)\n print(out.shape)\n print(dout.shape)\n\n target_out = out.new(*out.size())\n target_out.data.uniform_(-0.01, 0.01)\n target_dout = dout.new(*dout.size())\n target_dout.data.uniform_(-0.01, 0.01)\n e = (target_out - out).mean()\n e.backward()\n e = (target_dout - dout).mean()\n e.backward()\n\n\ndef example_mdpooling():\n input = torch.randn(2, 32, 64, 64).cuda()\n input.requires_grad = True\n batch_inds = torch.randint(2, (20, 1)).cuda().float()\n x = torch.randint(256, (20, 1)).cuda().float()\n y = torch.randint(256, (20, 1)).cuda().float()\n w = torch.randint(64, (20, 1)).cuda().float()\n h = torch.randint(64, (20, 1)).cuda().float()\n rois = torch.cat((batch_inds, x, y, x + w, y + h), dim=1)\n\n # mdformable pooling (V2)\n dpooling = DCNPooling(spatial_scale=1.0 / 4,\n pooled_size=7,\n output_dim=32,\n no_trans=False,\n group_size=1,\n trans_std=0.1,\n deform_fc_dim=1024).cuda()\n\n dout = dpooling(input, rois)\n target = dout.new(*dout.size())\n target.data.uniform_(-0.1, 0.1)\n error = (target - dout).mean()\n error.backward()\n print(dout.shape)\n\n\nif __name__ == '__main__':\n\n example_dconv()\n example_dpooling()\n example_mdpooling()\n\n check_pooling_zero_offset()\n # zero offset check\n if inC == outC:\n check_zero_offset()\n\n check_gradient_dpooling()\n check_gradient_dconv()\n # \"\"\"\n # ****** Note: backward is not reentrant error may not be a serious problem,\n # ****** since the max error is less than 1e-7,\n # ****** Still looking for what trigger this problem\n # \"\"\"\n"
},
{
"path": "code/real/bsrt/model/__init__.py",
"content": "import os\nfrom importlib import import_module\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.parallel as P\nimport torch.utils.model_zoo\nimport time\n\nclass Model(nn.Module):\n def __init__(self, args, ckp):\n super(Model, self).__init__()\n self.args = args\n if args.local_rank == 0:\n print(\"Making model: \", args.model)\n print(\"Patch size: \", args.patch_size)\n \n self.scale = args.scale\n self.idx_scale = 0\n self.input_large = (args.model == 'VDSR')\n self.self_ensemble = args.self_ensemble\n self.chop = args.chop\n self.precision = args.precision\n self.cpu = args.cpu\n self.device = torch.device('cpu' if args.cpu else 'cuda:%d' % args.local_rank)\n self.n_GPUs = args.n_GPUs\n self.save_models = args.save_models\n\n module = import_module('model.' + args.model.lower())\n self.model = module.make_model(args).to(self.device)\n\n if args.precision == 'half':\n self.model.half()\n\n self.load(\n ckp.get_path('model'),\n pre_train=args.pre_train,\n resume=args.resume,\n cpu=args.cpu\n )\n\n # time.sleep(3)\n\n if args.n_GPUs > 1:\n self.model = nn.parallel.DistributedDataParallel(self.model,\n device_ids=[args.local_rank],\n find_unused_parameters=True\n )\n\n print(self.model, file=ckp.log_file)\n\n def forward(self, x, idx_scale):\n self.idx_scale = idx_scale\n if hasattr(self.model, 'set_scale'):\n self.model.set_scale(idx_scale)\n\n if self.training:\n # if self.n_GPUs > 1:\n return self.model(x)\n else:\n if self.chop:\n forward_function = self.forward_chop\n else:\n forward_function = self.model.forward\n\n if self.self_ensemble:\n return self.forward_x8(x, forward_function=forward_function)\n else:\n # return self.model(x)\n return forward_function(x)\n\n def save(self, apath, epoch, is_best=False):\n save_dirs = [os.path.join(apath, 'model_latest.pt')]\n\n if is_best:\n save_dirs.append(os.path.join(apath, 'model_best.pt'))\n if self.save_models:\n save_dirs.append(\n os.path.join(apath, 'model_{}.pt'.format(epoch))\n )\n if self.n_GPUs > 1:\n model = self.model.module\n else:\n model = self.model\n\n for s in save_dirs:\n torch.save(self.model.state_dict(), s)\n\n def load(self, apath, pre_train='', resume=-1, cpu=False):\n load_from = None\n kwargs = {}\n if cpu:\n kwargs = {'map_location': lambda storage, loc: storage}\n\n if resume == -1:\n load_from = torch.load(\n os.path.join(apath, 'model_latest.pt'),\n **kwargs\n )\n elif resume == 0:\n if pre_train == 'download':\n print('Download the model')\n dir_model = os.path.join('..', 'models')\n os.makedirs(dir_model, exist_ok=True)\n load_from = torch.utils.model_zoo.load_url(\n self.model.url,\n model_dir=dir_model,\n **kwargs\n )\n elif pre_train:\n if self.args.local_rank == 0:\n print('Load the model from {}'.format(pre_train))\n map_location = {'cuda:%d' % 0: 'cuda:%d' % self.args.local_rank}\n load_from = torch.load(pre_train, map_location=map_location)\n # print(load_from.keys())\n else:\n load_from = torch.load(\n os.path.join(apath, 'model_{}.pt'.format(resume)),\n **kwargs\n )\n\n if load_from:\n self.model.load_state_dict(load_from, strict=True)\n del load_from\n\n \n if self.args.finetune:\n if self.args.local_rank == 0:\n print('finetune')\n for param in self.model.parameters():\n param.requires_grad = False\n\n for param in self.model.HRconv.parameters():\n param.requires_grad = True\n for param in self.model.conv_last.parameters():\n param.requires_grad = True\n\n if self.args.finetune_prelayer:\n if self.args.local_rank == 0:\n print('finetune_prelayer')\n if self.args.swinfeature:\n if self.args.model == 'MBSRT':\n for param in self.model.pre_layer1.parameters():\n param.requires_grad = True\n for param in self.model.pre_layer2.parameters():\n param.requires_grad = True\n else:\n for param in self.model.pre_layers.parameters():\n param.requires_grad = True\n else:\n for param in self.model.feature_extraction.parameters():\n param.requires_grad = True\n\n for param in self.model.conv_after_pre_layer.parameters():\n param.requires_grad = True\n\n if self.args.finetune_align:\n if self.args.local_rank == 0:\n print('finetune_align')\n for param in self.model.align.parameters():\n param.requires_grad = True\n\n if self.args.finetune_spynet:\n if self.args.local_rank == 0:\n print('finetune_spynet')\n for param in self.model.spynet.parameters():\n param.requires_grad = True\n\n if self.args.finetune_swin:\n if self.args.local_rank == 0:\n print('finetune_swin')\n for param in self.model.layers.parameters():\n param.requires_grad = True\n for param in self.model.conv_after_body.parameters():\n param.requires_grad = True\n\n if self.args.finetune_upconv:\n if self.args.local_rank == 0:\n print('finetune_upconv')\n for param in self.model.upconv1.parameters():\n param.requires_grad = True\n for param in self.model.upconv2.parameters():\n param.requires_grad = True\n for param in self.model.skipup1.parameters():\n param.requires_grad = True\n for param in self.model.skipup2.parameters():\n param.requires_grad = True\n\n if self.args.finetune_conv:\n if self.args.local_rank == 0:\n print('finetune_conv')\n # for param in self.model.conv_first.parameters():\n # param.requires_grad = True\n # for param in self.model.conv_flow.parameters():\n # param.requires_grad = True\n # for param in self.model.fea_L2_conv1.parameters():\n # param.requires_grad = True\n # for param in self.model.fea_L3_conv1.parameters():\n # param.requires_grad = True\n # for param in self.model.toplayer.parameters():\n # param.requires_grad = True\n # for param in self.model.smooth1.parameters():\n # param.requires_grad = True\n # for param in self.model.smooth2.parameters():\n # param.requires_grad = True\n # for param in self.model.latlayer1.parameters():\n # param.requires_grad = True\n # for param in self.model.latlayer2.parameters():\n # param.requires_grad = True\n # for param in self.model.fusion.parameters():\n # param.requires_grad = True\n # for param in self.model.conv_after_pre_layer.parameters():\n # param.requires_grad = True\n for param in self.model.conv_after_body.parameters():\n param.requires_grad = True\n \n \n\n def forward_chop(self, *args, shave=10, min_size=160000):\n scale = 1 if self.input_large else self.scale[self.idx_scale]\n n_GPUs = min(self.n_GPUs, 4)\n # height, width\n h, w = args[0].size()[-2:]\n\n top = slice(0, h//2 + shave)\n bottom = slice(h - h//2 - shave, h)\n left = slice(0, w//2 + shave)\n right = slice(w - w//2 - shave, w)\n x_chops = [torch.cat([\n a[..., top, left],\n a[..., top, right],\n a[..., bottom, left],\n a[..., bottom, right]\n ]) for a in args]\n\n y_chops = []\n if h * w < 4 * min_size:\n for i in range(0, 4, n_GPUs):\n x = [x_chop[i:(i + n_GPUs)] for x_chop in x_chops]\n y = P.data_parallel(self.model, *x, range(n_GPUs))\n if not isinstance(y, list): y = [y]\n if not y_chops:\n y_chops = [[c for c in _y.chunk(n_GPUs, dim=0)] for _y in y]\n else:\n for y_chop, _y in zip(y_chops, y):\n y_chop.extend(_y.chunk(n_GPUs, dim=0))\n else:\n for p in zip(*x_chops):\n y = self.forward_chop(*p, shave=shave, min_size=min_size)\n if not isinstance(y, list): y = [y]\n if not y_chops:\n y_chops = [[_y] for _y in y]\n else:\n for y_chop, _y in zip(y_chops, y): y_chop.append(_y)\n\n h *= scale\n w *= scale\n top = slice(0, h//2)\n bottom = slice(h - h//2, h)\n bottom_r = slice(h//2 - h, None)\n left = slice(0, w//2)\n right = slice(w - w//2, w)\n right_r = slice(w//2 - w, None)\n\n # batch size, number of color channels\n b, c = y_chops[0][0].size()[:-2]\n y = [y_chop[0].new(b, c, h, w) for y_chop in y_chops]\n for y_chop, _y in zip(y_chops, y):\n _y[..., top, left] = y_chop[0][..., top, left]\n _y[..., top, right] = y_chop[1][..., top, right_r]\n _y[..., bottom, left] = y_chop[2][..., bottom_r, left]\n _y[..., bottom, right] = y_chop[3][..., bottom_r, right_r]\n\n if len(y) == 1: y = y[0]\n\n return y\n\n def forward_x8(self, *args, forward_function=None):\n def _transform(v, op):\n if self.precision != 'single': v = v.float()\n\n v2np = v.data.cpu().numpy()\n if op == 'v':\n tfnp = v2np[:, :, :, ::-1].copy()\n elif op == 'h':\n tfnp = v2np[:, :, ::-1, :].copy()\n elif op == 't':\n tfnp = v2np.transpose((0, 1, 3, 2)).copy()\n\n ret = torch.Tensor(tfnp).to(self.device)\n if self.precision == 'half': ret = ret.half()\n\n return ret\n\n list_x = []\n for a in args:\n x = [a]\n for tf in 'v', 'h', 't': x.extend([_transform(_x, tf) for _x in x])\n\n list_x.append(x)\n\n list_y = []\n for x in zip(*list_x):\n y = forward_function(*x)\n if not isinstance(y, list): y = [y]\n if not list_y:\n list_y = [[_y] for _y in y]\n else:\n for _list_y, _y in zip(list_y, y): _list_y.append(_y)\n\n for _list_y in list_y:\n for i in range(len(_list_y)):\n if i > 3:\n _list_y[i] = _transform(_list_y[i], 't')\n if i % 4 > 1:\n _list_y[i] = _transform(_list_y[i], 'h')\n if (i % 4) % 2 == 1:\n _list_y[i] = _transform(_list_y[i], 'v')\n\n y = [torch.cat(_y, dim=0).mean(dim=0, keepdim=True) for _y in list_y]\n if len(y) == 1: y = y[0]\n\n return y\n"
},
{
"path": "code/real/bsrt/model/arch_util.py",
"content": "import torch\nimport torch.nn as nn\nimport torch.nn.init as init\nimport torch.nn.functional as F\nfrom model import common\nfrom model.utils.psconv import PSGConv2d as PSConv2d, PyConv2d\n\n\ndef initialize_weights(net_l, scale=1):\n if not isinstance(net_l, list):\n net_l = [net_l]\n for net in net_l:\n for m in net.modules():\n if isinstance(m, nn.Conv2d):\n init.kaiming_normal_(m.weight, a=0, mode='fan_in')\n m.weight.data *= scale # for residual block\n if m.bias is not None:\n m.bias.data.zero_()\n elif isinstance(m, nn.Linear):\n init.kaiming_normal_(m.weight, a=0, mode='fan_in')\n m.weight.data *= scale\n if m.bias is not None:\n m.bias.data.zero_()\n elif isinstance(m, nn.BatchNorm2d):\n init.constant_(m.weight, 1)\n init.constant_(m.bias.data, 0.0)\n\n\ndef make_layer(block, n_layers):\n layers = []\n for _ in range(n_layers):\n layers.append(block())\n return nn.Sequential(*layers)\n\n\n###########################\n\ndef conv_layer(in_channels, out_channels, kernel_size, stride=1, padding=0):\n return nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, bias=True)\n \n\nclass ESA(nn.Module):\n def __init__(self, n_feats, conv=conv_layer):\n super(ESA, self).__init__()\n f = n_feats // 4\n self.conv1 = conv(n_feats, f, kernel_size=1)\n self.conv_f = conv(f, f, kernel_size=1)\n self.conv_max = conv(f, f, kernel_size=3, padding=1)\n self.conv2 = conv(f, f, kernel_size=3, stride=2, padding=0)\n self.conv3 = conv(f, f, kernel_size=3, padding=1)\n self.conv3_ = conv(f, f, kernel_size=3, padding=1)\n self.conv4 = conv(f, n_feats, kernel_size=1)\n self.sigmoid = nn.Sigmoid()\n self.relu = nn.ReLU(inplace=True)\n\n def forward(self, x):\n c1_ = (self.conv1(x))\n c1 = self.conv2(c1_)\n v_max = F.max_pool2d(c1, kernel_size=7, stride=3)\n v_range = self.relu(self.conv_max(v_max))\n c3 = self.relu(self.conv3(v_range))\n c3 = self.conv3_(c3)\n c3 = F.interpolate(c3, (x.size(2), x.size(3)), mode='bilinear', align_corners=False) \n cf = self.conv_f(c1_)\n c4 = self.conv4(c3+cf)\n m = self.sigmoid(c4)\n \n return x * m\n\n\nclass DWConv(nn.Module):\n def __init__(self, dim=768):\n super(DWConv, self).__init__()\n self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)\n\n def forward(self, x):\n x = self.dwconv(x)\n return x\n\n##########################\n\nclass SELayer(nn.Module):\n '''\n SE-block\n '''\n def __init__(self, channel, reduction=16):\n super(SELayer, self).__init__()\n self.avg_pool = nn.AdaptiveAvgPool2d(1)\n self.fc = nn.Sequential(\n nn.Linear(channel, channel // reduction, bias=False),\n nn.ReLU(inplace=True),\n nn.Linear(channel // reduction, channel, bias=False),\n # nn.Sigmoid()\n )\n\n def forward(self, x):\n b, c, _, _ = x.size()\n y = self.avg_pool(x).view(b, c)\n y = self.fc(y).view(b, c, 1, 1)\n return x * y.expand_as(x)\n\nclass ResidualBlock_noBN(nn.Module):\n '''Residual block w/o BN\n ---Conv-ReLU-Conv-+-\n |________________|\n '''\n\n def __init__(self, nf=64):\n super(ResidualBlock_noBN, self).__init__()\n self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n\n # initialization\n initialize_weights([self.conv1, self.conv2], 0.1)\n\n def forward(self, x):\n identity = x\n out = F.relu(self.conv1(x), inplace=True)\n out = self.conv2(out)\n return identity + out\n\n\nclass ResidualBlock_SE(nn.Module):\n '''Residual block w/o BN\n ---Conv-ReLU-Conv-+-\n |________________|\n '''\n\n def __init__(self, nf=64, reduction=16):\n super(ResidualBlock_SE, self).__init__()\n self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.conv3 = nn.Conv2d(3 * nf, nf, 1, padding=0, dilation=1, bias=True)\n self.se = SELayer(nf, reduction)\n # initialization\n initialize_weights([self.conv1, self.conv2, self.conv3], 0.1)\n\n def forward(self, x):\n identity = x\n basic_out = F.relu(self.conv1(x), inplace=True)\n basic_out = self.conv2(basic_out)\n se_out = self.se(basic_out)\n out = torch.cat((identity, basic_out, se_out), 1)\n out = self.conv3(out)\n return out\n\n\nclass _PositionAttentionModule(nn.Module):\n \"\"\" Position attention module\"\"\"\n\n def __init__(self, in_channels, **kwargs):\n super(_PositionAttentionModule, self).__init__()\n self.conv_b = nn.Conv2d(in_channels, in_channels // 8, 1)\n self.conv_c = nn.Conv2d(in_channels, in_channels // 8, 1)\n self.conv_d = nn.Conv2d(in_channels, in_channels, 1)\n self.alpha = nn.Parameter(torch.zeros(1))\n self.softmax = nn.Softmax(dim=-1)\n\n def forward(self, x):\n batch_size, _, height, width = x.size()\n feat_b = self.conv_b(x).view(batch_size, -1, height * width).permute(0, 2, 1)\n feat_c = self.conv_c(x).view(batch_size, -1, height * width)\n attention_s = self.softmax(torch.bmm(feat_b, feat_c))\n feat_d = self.conv_d(x).view(batch_size, -1, height * width)\n feat_e = torch.bmm(feat_d, attention_s.permute(0, 2, 1)).view(batch_size, -1, height, width)\n out = self.alpha * feat_e + x\n\n return out\n\n## Spatial Attention (CA) Layer\nclass SALayer(nn.Module):\n def __init__(self, wn=None):\n super(SALayer,self).__init__()\n self.body = nn.Sequential(\n wn(nn.Conv2d(2, 1, 7, 1, 3, bias=False)),\n nn.Sigmoid()\n )\n def forward(self, x):\n avg_f = torch.mean(x, dim=1, keepdim=True)\n max_f = torch.max(x, dim=1, keepdim=True)[0]\n y = torch.cat([avg_f, max_f], dim=1)\n return self.body(y).expand_as(x) * x\n\n\n## Channel Attention (CA) Layer\nclass CALayerV2(nn.Module):\n def __init__(self, n_feat, reduction=16, wn=None):\n super(CALayerV2, self).__init__()\n # global average pooling: feature --> point\n self.avg_pool = nn.AdaptiveAvgPool2d(1)\n self.max_pool = nn.AdaptiveMaxPool2d(1)\n # feature channel downscale and upscale --> channel weight\n self.conv_du = nn.Sequential(\n wn(nn.Conv2d(n_feat, n_feat//reduction, 1, padding=0, bias=False)),\n nn.ReLU(inplace=True),\n wn(nn.Conv2d(n_feat//reduction, n_feat, 1, padding=0, bias=False)),\n # nn.Sigmoid()\n )\n\n def forward(self, x):\n y1 = self.avg_pool(x)\n y2 = self.max_pool(x)\n y1 = self.conv_du(y1)\n y2 = self.conv_du(y2)\n return x * torch.sigmoid(y1+y2)\n\nclass DALayer(nn.Module):\n def __init__(self, channel, reduction, wn):\n super(DALayer, self).__init__()\n # global average pooling: feature --> point\n self.ca = CALayer(channel, reduction, wn)\n self.sa = SALayer(wn)\n self.conv = wn(nn.Conv2d(channel*2, channel, 1))\n\n def forward(self, x):\n ca = self.ca(x)\n sa = self.sa(x)\n res = self.conv(torch.cat([ca, sa], dim=1))\n return res + x\n\n\n## Channel Attention (CA) Layer\nclass CALayer(nn.Module):\n def __init__(self, channel, reduction, wn):\n super(CALayer, self).__init__()\n # global average pooling: feature --> point\n self.avg_pool = nn.AdaptiveAvgPool2d(1)\n # feature channel downscale and upscale --> channel weight\n self.conv_du = nn.Sequential(\n wn(nn.Conv2d(channel, channel // reduction, 1, padding=0, bias=True)),\n nn.ReLU(inplace=True),\n wn(nn.Conv2d(channel // reduction, channel, 1, padding=0, bias=True)),\n nn.Sigmoid()\n )\n\n def forward(self, x):\n y = self.avg_pool(x)\n y = self.conv_du(y)\n return x * y\n\n\n## Residual Channel Attention Block (RCAB)\nclass RCAB(nn.Module):\n def __init__(\n self, conv, n_feat, kernel_size, reduction, wn,\n bias=True, bn=False, act=nn.ReLU(True), res_scale=1, da=False):\n\n super(RCAB, self).__init__()\n\n expand = 6\n linear = 0.75\n modules_body = []\n # for i in range(2):\n modules_body.append(wn(nn.Conv2d(n_feat, n_feat*expand, 1, bias=bias)))\n modules_body.append(act)\n modules_body.append(wn(nn.Conv2d(n_feat*expand, int(n_feat*linear), 1, bias=bias)))\n modules_body.append(conv(int(n_feat*linear), n_feat, kernel_size, bias=bias))\n if da:\n modules_body.append(DALayer(n_feat, reduction, wn))\n else:\n modules_body.append(CALayer(n_feat, reduction, wn))\n\n self.body = nn.Sequential(*modules_body)\n self.res_scale = res_scale\n\n def forward(self, x):\n res = self.body(x)\n #res = self.body(x).mul(self.res_scale)\n res += x\n return res\n\n## Residual Group (RG)\nclass ResidualGroup(nn.Module):\n def __init__(self, n_feat, n_resblocks, da=False):\n super(ResidualGroup, self).__init__()\n kernel_size = 3\n res_scale = 1\n reduction = 16\n\n conv = common.default_conv\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_body = []\n modules_body = [\n RCAB(\n conv, n_feat, kernel_size, reduction, wn=wn, bias=True,\n bn=False, act=nn.ReLU(True), res_scale=res_scale, da=da) \\\n for _ in range(n_resblocks)]\n modules_body.append(wn(conv(n_feat, n_feat, kernel_size)))\n self.body = nn.Sequential(*modules_body)\n\n def forward(self, x):\n res = self.body(x)\n res += x\n return res\n\n\n################################################################\n################################################################\n################################################################\n\ndef make_layer_idx(block, n_layers):\n layers = []\n for i in range(n_layers):\n layers.append(block(idx=i))\n return nn.Sequential(*layers)\n\n## Residual Channel Attention Block (RCAB)\nclass LRSCRCAB(nn.Module):\n def __init__(\n self, conv, n_feat, kernel_size, reduction, wn,\n bias=True, bn=False, act=nn.ReLU(True), res_scale=1, da=False, idx=0):\n super(LRSCRCAB, self).__init__()\n\n expand = 6\n linear = 0.75\n\n modules_body = [wn(nn.Conv2d(n_feat*(idx+1), n_feat, 1, 1, 0, bias=True))] if idx > 0 else []\n # for i in range(2):\n modules_body.append(wn(nn.Conv2d(n_feat, n_feat*expand, 1, bias=bias)))\n modules_body.append(act)\n modules_body.append(wn(nn.Conv2d(n_feat*expand, int(n_feat*linear), 1, bias=bias)))\n modules_body.append(wn(conv(int(n_feat*linear), n_feat, kernel_size, bias=bias)))\n if da:\n modules_body.append(DALayer(n_feat, reduction, wn))\n else:\n modules_body.append(CALayer(n_feat, reduction, wn))\n\n self.body = nn.Sequential(*modules_body)\n self.res_scale = res_scale\n\n def forward(self, x):\n res = self.body(x)\n res = torch.cat([res, x], dim=1)\n return res\n\n\n## Residual Channel Attention Block (RCAB)\nclass LRSCPYRCAB(nn.Module):\n def __init__(\n self, conv, n_feat, kernel_size, reduction, wn,\n bias=True, bn=False, act=nn.ReLU(True), res_scale=1, da=False, idx=0):\n super(LRSCPYRCAB, self).__init__()\n\n expand = 6\n linear = 0.75\n\n modules_body = [wn(nn.Conv2d(n_feat*(idx+1), n_feat, 1, 1, 0, bias=True))] if idx > 0 else []\n # for i in range(2):\n modules_body.append(wn(nn.Conv2d(n_feat, n_feat*expand, 1, bias=bias)))\n modules_body.append(act)\n modules_body.append(wn(nn.Conv2d(n_feat*expand, int(n_feat*linear), 1, bias=bias)))\n modules_body.append(\n PyConv2d(in_channels=int(n_feat*linear),\n out_channels=[n_feat//4, n_feat//4, n_feat//2],\n pyconv_kernels=[3, 5, 7],\n pyconv_groups=[1, 4, 8]))\n if da:\n modules_body.append(DALayer(n_feat, reduction, wn))\n else:\n modules_body.append(CALayer(n_feat, reduction, wn))\n\n self.body = nn.Sequential(*modules_body)\n self.res_scale = res_scale\n\n def forward(self, x):\n res = self.body(x)\n res = torch.cat([res, x], dim=1)\n return res\n\n## Long-Range Skip-connect Residual Group (RG)\nclass LRSCResidualGroup(nn.Module):\n def __init__(self, n_feat, n_resblocks, da=False, idx=0):\n super(LRSCResidualGroup, self).__init__()\n kernel_size = 3\n res_scale = 1\n reduction = 16\n\n conv = common.default_conv\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_head = [wn(conv(n_feat*(idx+1), n_feat, 1, bias=True))] if idx > 0 else []\n modules_body = [\n LRSCRCAB(\n conv, n_feat, kernel_size, reduction, wn=wn, bias=True,\n bn=False, act=nn.ReLU(True), res_scale=res_scale, da=da, idx=i) \\\n for i in range(n_resblocks)]\n modules_body.append(wn(conv(n_feat*(n_resblocks+1), n_feat, kernel_size)))\n self.head = nn.Sequential(*modules_head)\n self.body = nn.Sequential(*modules_body)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = torch.cat([res, x], dim=1)\n return res\n\n\n## Long-Range Skip-connect Residual Group (RG)\nclass LRSCPSResidualGroup(nn.Module):\n def __init__(self, n_feat, n_resblocks, da=False, idx=0):\n super(LRSCPSResidualGroup, self).__init__()\n kernel_size = 3\n res_scale = 1\n reduction = 16\n\n conv = PSConv2d\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_head = [wn(nn.Conv2d(n_feat*(idx+1), n_feat, 1, 1, 0, bias=True))] if idx > 0 else []\n modules_body = [\n LRSCRCAB(\n conv, n_feat, kernel_size, reduction, wn=wn, bias=True,\n bn=False, act=nn.ReLU(True), res_scale=res_scale, da=da, idx=i) \\\n for i in range(n_resblocks)]\n modules_tail = [wn(conv(n_feat*(n_resblocks+1), n_feat, kernel_size))]\n self.head = nn.Sequential(*modules_head)\n self.body = nn.Sequential(*modules_body)\n self.tail = nn.Sequential(*modules_tail)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = self.tail(res)\n res = torch.cat([res, x], dim=1)\n return res\n\n\n## Long-Range Skip-connect Residual Group (RG)\nclass LRSCPyResidualGroup(nn.Module):\n def __init__(self, n_feat, n_resblocks, da=False, idx=0):\n super(LRSCPyResidualGroup, self).__init__()\n kernel_size = 3\n res_scale = 1\n reduction = 16\n\n conv = PyConv2d\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_head = [wn(nn.Conv2d(n_feat*(idx+1), n_feat, 1, 1, 0, bias=True))] if idx > 0 else []\n modules_body = [\n LRSCPYRCAB(\n conv, n_feat, kernel_size, reduction, wn=wn, bias=True,\n bn=False, act=nn.ReLU(True), res_scale=res_scale, da=da, idx=i) \\\n for i in range(n_resblocks)]\n modules_tail = [wn(nn.Conv2d(n_feat*(n_resblocks+1), n_feat, 1))]\n self.head = nn.Sequential(*modules_head)\n self.body = nn.Sequential(*modules_body)\n self.tail = nn.Sequential(*modules_tail)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = self.tail(res)\n res = torch.cat([res, x], dim=1)\n return res\n\nclass LRSCWideActResBlock(nn.Module):\n def __init__(self, nf=64, idx=0):\n super(LRSCWideActResBlock, self).__init__()\n self.res_scale = 1\n\n expand = 6\n linear = 0.8\n kernel_size = 3\n wn = lambda x: torch.nn.utils.weight_norm(x)\n act=nn.ReLU(True)\n head = [wn(nn.Conv2d(nf*(idx+1), nf, 1, bias=True))] if idx > 0 else []\n\n body = []\n body.append(\n wn(nn.Conv2d(nf, nf*expand, 1, padding=1//2)))\n body.append(act)\n body.append(\n wn(nn.Conv2d(nf*expand, int(nf*linear), 1, padding=1//2)))\n body.append(\n wn(nn.Conv2d(int(nf*linear), nf, kernel_size, padding=kernel_size//2)))\n\n self.head = nn.Sequential(*head)\n self.body = nn.Sequential(*body)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = torch.cat([res, x], dim=1)\n return res\n\nclass LRSCPyWideActResBlock(nn.Module):\n def __init__(self, nf=64, idx=0):\n super(LRSCPyWideActResBlock, self).__init__()\n self.res_scale = 1\n\n expand = 6\n linear = 0.75\n kernel_size = 3\n wn = lambda x: torch.nn.utils.weight_norm(x)\n act=nn.ReLU(True)\n head = [wn(nn.Conv2d(nf*(idx+1), nf, 1, bias=True))] if idx > 0 else []\n\n body = []\n body.append(\n wn(nn.Conv2d(nf, nf*expand, 1, padding=1//2)))\n body.append(act)\n body.append(\n wn(nn.Conv2d(nf*expand, int(nf*linear), 1, padding=1//2)))\n body.append(\n PyConv2d(in_channels=int(nf*linear),\n out_channels=[nf//4, nf//4, nf//2],\n pyconv_kernels=[3, 5, 7],\n pyconv_groups=[1, 4, 8]))\n\n self.head = nn.Sequential(*head)\n self.body = nn.Sequential(*body)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = torch.cat([res, x], dim=1)\n return res\n\n\n## Long-Range Skip-connect Residual Group (RG)\nclass LRSCPyWideActResGroup(nn.Module):\n def __init__(self, nf, n_resblocks, idx=0):\n super(LRSCPyWideActResGroup, self).__init__()\n kernel_size = 3\n\n conv = PyConv2d\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_head = [wn(nn.Conv2d(nf*(idx+1), nf, 1, 1, 0, bias=True))] if idx > 0 else []\n modules_body = [\n LRSCPyWideActResBlock(nf=nf, idx=i) for i in range(n_resblocks)]\n modules_tail = [wn(nn.Conv2d(nf*(n_resblocks+1), nf, 1))]\n self.head = nn.Sequential(*modules_head)\n self.body = nn.Sequential(*modules_body)\n self.tail = nn.Sequential(*modules_tail)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = self.tail(res)\n res = torch.cat([res, x], dim=1)\n return res\n\n\n## Long-Range Skip-connect Residual Group (RG)\nclass LRSCWideActResGroup(nn.Module):\n def __init__(self, nf, n_resblocks, idx=0):\n super(LRSCWideActResGroup, self).__init__()\n kernel_size = 3\n\n conv = PyConv2d\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_head = [wn(nn.Conv2d(nf*(idx+1), nf, 1, 1, 0, bias=True))] if idx > 0 else []\n modules_body = [\n LRSCWideActResBlock(nf=nf, idx=i) for i in range(n_resblocks)]\n modules_tail = [wn(nn.Conv2d(nf*(n_resblocks+1), nf, 1))]\n self.head = nn.Sequential(*modules_head)\n self.body = nn.Sequential(*modules_body)\n self.tail = nn.Sequential(*modules_tail)\n\n def forward(self, x):\n res = self.head(x)\n res = self.body(res)\n res = self.tail(res)\n res = torch.cat([res, x], dim=1)\n return res\n\n################################################################\n################################################################\n################################################################\n\n\n## Residual Channel Attention Block (RCAB)\nclass PYRCAB(nn.Module):\n def __init__(\n self, conv, n_feat, kernel_size, reduction, wn,\n bias=True, bn=False, act=nn.ReLU(True), res_scale=1, da=False):\n super(PYRCAB, self).__init__()\n\n expand = 6\n linear = 0.75\n modules_body = []\n # for i in range(2):\n modules_body.append(wn(nn.Conv2d(n_feat, n_feat*expand, 1, bias=bias)))\n modules_body.append(act)\n modules_body.append(wn(nn.Conv2d(n_feat*expand, int(n_feat*linear), 1, bias=bias)))\n # modules_body.append(conv(, n_feat, kernel_size, bias=bias))\n modules_body.append(PyConv2d(in_channels=int(n_feat*linear),\n out_channels=[n_feat//4, n_feat//4, n_feat//2],\n pyconv_kernels=[3, 5, 7],\n pyconv_groups=[1, 4, 8], bias=bias))\n if da:\n modules_body.append(DALayer(n_feat, reduction, wn))\n else:\n modules_body.append(CALayer(n_feat, reduction, wn))\n\n self.body = nn.Sequential(*modules_body)\n self.res_scale = res_scale\n\n def forward(self, x):\n res = self.body(x)\n res += x\n return res\n\n## Residual Group (RG)\nclass PyResidualGroup(nn.Module):\n def __init__(self, n_feat, n_resblocks, da=False):\n super(PyResidualGroup, self).__init__()\n kernel_size = 3\n res_scale = 1\n reduction = 16\n\n conv = PyConv2d\n wn = lambda x: torch.nn.utils.weight_norm(x)\n\n modules_body = []\n modules_body = [\n PYRCAB(\n conv, n_feat, kernel_size, reduction, wn=wn, bias=True,\n bn=False, act=nn.ReLU(True), res_scale=res_scale, da=da) \\\n for _ in range(n_resblocks)]\n modules_body.append(\n PyConv2d(in_channels=n_feat,\n out_channels=[n_feat//4, n_feat//4, n_feat//2],\n pyconv_kernels=[3, 5, 7],\n pyconv_groups=[1, 4, 8]))\n self.body = nn.Sequential(*modules_body)\n\n def forward(self, x):\n res = self.body(x)\n res += x\n return res\n\nclass WideActResBlock(nn.Module):\n def __init__(self, nf=64):\n super(WideActResBlock, self).__init__()\n self.res_scale = 1\n body = []\n expand = 6\n linear = 0.8\n kernel_size = 3\n wn = lambda x: torch.nn.utils.weight_norm(x)\n act=nn.ReLU(True)\n\n body.append(\n wn(nn.Conv2d(nf, nf*expand, 1, padding=1//2)))\n body.append(act)\n body.append(\n wn(nn.Conv2d(nf*expand, int(nf*linear), 1, padding=1//2)))\n body.append(\n wn(nn.Conv2d(int(nf*linear), nf, kernel_size, padding=kernel_size//2)))\n\n self.body = nn.Sequential(*body)\n\n def forward(self, x):\n res = self.body(x) * self.res_scale\n res += x\n return res\n\n\nclass PSWideActResBlock(nn.Module):\n def __init__(self, nf=64):\n super(PSWideActResBlock, self).__init__()\n self.res_scale = 1\n body = []\n expand = 6\n linear = 0.75\n kernel_size = 3\n wn = lambda x: torch.nn.utils.weight_norm(x)\n act=nn.ReLU(True)\n\n body.append(\n wn(nn.Conv2d(nf, nf*expand, 1, padding=1//2)))\n body.append(act)\n body.append(\n wn(nn.Conv2d(nf*expand, int(nf*linear), 1, padding=1//2)))\n body.append(\n wn(PSConv2d(int(nf*linear), nf, kernel_size, padding=kernel_size//2)))\n\n self.body = nn.Sequential(*body)\n\n def forward(self, x):\n res = self.body(x) * self.res_scale\n res += x\n return res\n\n\nclass PyWideActResBlock(nn.Module):\n def __init__(self, nf=64):\n super(PyWideActResBlock, self).__init__()\n self.res_scale = 1\n body = []\n expand = 6\n linear = 0.75\n kernel_size = 3\n wn = lambda x: torch.nn.utils.weight_norm(x)\n act=nn.ReLU(True)\n expand_nf = nf*expand\n linear_nf = int(nf * linear)\n\n body.append(\n wn(nn.Conv2d(nf, nf*expand, 1, padding=1//2)))\n body.append(act)\n body.append(\n wn(nn.Conv2d(nf*expand, int(nf*linear), 1, padding=1//2)))\n body.append(\n PyConv2d(in_channels=linear_nf,\n out_channels=[nf//4, nf//4, nf//2],\n pyconv_kernels=[3, 5, 7],\n pyconv_groups=[1, 4, 8]))\n\n self.body = nn.Sequential(*body)\n\n def forward(self, x):\n res = self.body(x) * self.res_scale\n res += x\n return res\n\n\ndef flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros', align_corners=True, use_pad_mask=False):\n \"\"\"Warp an image or feature map with optical flow.\n\n Args:\n x (Tensor): Tensor with size (n, c, h, w).\n flow (Tensor): Tensor with size (n, h, w, 2), normal value.\n interp_mode (str): 'nearest' or 'bilinear' or 'nearest4'. Default: 'bilinear'.\n padding_mode (str): 'zeros' or 'border' or 'reflection'.\n Default: 'zeros'.\n align_corners (bool): Before pytorch 1.3, the default value is\n align_corners=True. After pytorch 1.3, the default value is\n align_corners=False. Here, we use the True as default.\n use_pad_mask (bool): only used for PWCNet, x is first padded with ones along the channel dimension.\n The mask is generated according to the grid_sample results of the padded dimension.\n\n\n Returns:\n Tensor: Warped image or feature map.\n \"\"\"\n # assert x.size()[-2:] == flow.size()[1:3] # temporaily turned off for image-wise shift\n n, _, h, w = x.size()\n x = x.float()\n # create mesh grid\n # grid_y, grid_x = torch.meshgrid(torch.arange(0, h).type_as(x), torch.arange(0, w).type_as(x)) # an illegal memory access on TITAN RTX + PyTorch1.9.1\n grid_y, grid_x = torch.meshgrid(torch.arange(0, h, dtype=x.dtype, device=x.device), torch.arange(0, w, dtype=x.dtype, device=x.device))\n grid = torch.stack((grid_x, grid_y), 2).float() # W(x), H(y), 2\n grid.requires_grad = False\n grid = grid.type_as(x)\n vgrid = grid + flow\n\n # if use_pad_mask: # for PWCNet\n # x = F.pad(x, (0,0,0,0,0,1), mode='constant', value=1)\n\n # scale grid to [-1,1]\n if interp_mode == 'nearest4': # todo: bug, no gradient for flow model in this case!!! but the result is good\n vgrid_x_floor = 2.0 * torch.floor(vgrid[:, :, :, 0]) / max(w - 1, 1) - 1.0\n vgrid_x_ceil = 2.0 * torch.ceil(vgrid[:, :, :, 0]) / max(w - 1, 1) - 1.0\n vgrid_y_floor = 2.0 * torch.floor(vgrid[:, :, :, 1]) / max(h - 1, 1) - 1.0\n vgrid_y_ceil = 2.0 * torch.ceil(vgrid[:, :, :, 1]) / max(h - 1, 1) - 1.0\n\n output00 = F.grid_sample(x, torch.stack((vgrid_x_floor, vgrid_y_floor), dim=3), mode='nearest', padding_mode=padding_mode, align_corners=align_corners)\n output01 = F.grid_sample(x, torch.stack((vgrid_x_floor, vgrid_y_ceil), dim=3), mode='nearest', padding_mode=padding_mode, align_corners=align_corners)\n output10 = F.grid_sample(x, torch.stack((vgrid_x_ceil, vgrid_y_floor), dim=3), mode='nearest', padding_mode=padding_mode, align_corners=align_corners)\n output11 = F.grid_sample(x, torch.stack((vgrid_x_ceil, vgrid_y_ceil), dim=3), mode='nearest', padding_mode=padding_mode, align_corners=align_corners)\n\n return torch.cat([output00, output01, output10, output11], 1)\n\n else:\n vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(w - 1, 1) - 1.0\n vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(h - 1, 1) - 1.0\n vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)\n output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode, align_corners=align_corners)\n\n # if use_pad_mask: # for PWCNet\n # output = _flow_warp_masking(output)\n\n # TODO, what if align_corners=False\n return output\n\n\n# def flow_warp(x, flow, interp_mode='bilinear', padding_mode='zeros'):\n# \"\"\"Warp an image or feature map with optical flow\n# Args:\n# x (Tensor): size (N, C, H, W)\n# flow (Tensor): size (N, H, W, 2), normal value\n# interp_mode (str): 'nearest' or 'bilinear'\n# padding_mode (str): 'zeros' or 'border' or 'reflection'\n\n# Returns:\n# Tensor: warped image or feature map\n# \"\"\"\n# assert x.size()[-2:] == flow.size()[1:3]\n# B, C, H, W = x.size()\n# # mesh grid\n# grid_y, grid_x = torch.meshgrid(torch.arange(0, H), torch.arange(0, W))\n# grid = torch.stack((grid_x, grid_y), 2).float() # W(x), H(y), 2\n# grid.requires_grad = False\n# grid = grid.type_as(x)\n# vgrid = grid + flow\n# # scale grid to [-1,1]\n# vgrid_x = 2.0 * vgrid[:, :, :, 0] / max(W - 1, 1) - 1.0\n# vgrid_y = 2.0 * vgrid[:, :, :, 1] / max(H - 1, 1) - 1.0\n# vgrid_scaled = torch.stack((vgrid_x, vgrid_y), dim=3)\n# output = F.grid_sample(x, vgrid_scaled, mode=interp_mode, padding_mode=padding_mode)\n# return output\n"
},
{
"path": "code/real/bsrt/model/bsrt.py",
"content": "import functools\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport model.arch_util as arch_util\nfrom torch.cuda.amp import autocast\nimport model.swin_util as swu\nimport time\nimport os\nimport math\nfrom utils.debayer import Debayer3x3\nimport torchvision.utils as tvutils\nfrom datasets.burstsr_dataset import pack_raw_image, flatten_raw_image_batch\n\ntry:\n from model.non_local.non_local_cross_dot_product import NONLocalBlock2D as NonLocalCross\n from model.non_local.non_local_dot_product import NONLocalBlock2D as NonLocal\nexcept ImportError:\n raise ImportError('Failed to import Non_Local module.')\n\ntry:\n from model.DCNv2.dcn_v2 import DCN_sep as DCN, FlowGuidedDCN, InsideFlowGuidedDCN\nexcept ImportError:\n raise ImportError('Failed to import DCNv2 module.')\n\n\ndef make_model(args, parent=False):\n nframes = args.burst_size\n img_size = args.patch_size * 2\n patch_size = 1\n in_chans = args.burst_channel\n out_chans = args.n_colors\n \n if args.model_level == \"S\":\n depths = [6]*1 + [6] * 4\n num_heads = [6]*1 + [6] * 4\n embed_dim = 60\n elif args.model_level == \"L\":\n depths = [6]*1 + [8] * 6\n num_heads = [6]*1 + [6] * 6\n embed_dim = 180\n window_size = 8\n mlp_ratio = 2\n upscale = args.scale[0]\n non_local = args.non_local\n use_checkpoint=args.use_checkpoint\n\n if args.local_rank <= 0:\n print(\"depths: \", depths)\n\n return BSRT(args=args,nframes=nframes,\n img_size=img_size,\n patch_size=patch_size,\n in_chans=in_chans,\n out_chans=out_chans,\n embed_dim=embed_dim,\n depths=depths,\n num_heads=num_heads,\n window_size=window_size,\n mlp_ratio=mlp_ratio,\n upscale=upscale,\n non_local=non_local,\n use_checkpoint=use_checkpoint)\n\n\nclass BasicModule(nn.Module):\n \"\"\"Basic Module for SpyNet.\n \"\"\"\n\n def __init__(self):\n super(BasicModule, self).__init__()\n\n self.basic_module = nn.Sequential(\n nn.Conv2d(in_channels=8, out_channels=32, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),\n nn.Conv2d(in_channels=32, out_channels=64, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),\n nn.Conv2d(in_channels=64, out_channels=32, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),\n nn.Conv2d(in_channels=32, out_channels=16, kernel_size=7, stride=1, padding=3), nn.ReLU(inplace=False),\n nn.Conv2d(in_channels=16, out_channels=2, kernel_size=7, stride=1, padding=3))\n\n def forward(self, tensor_input):\n return self.basic_module(tensor_input)\n\n\nclass SpyNet(nn.Module):\n \"\"\"SpyNet architecture.\n\n Args:\n load_path (str): path for pretrained SpyNet. Default: None.\n return_levels (list[int]): return flows of different levels. Default: [5].\n \"\"\"\n\n def __init__(self, load_path=None, return_levels=[5]):\n super(SpyNet, self).__init__()\n self.return_levels = return_levels\n self.basic_module = nn.ModuleList([BasicModule() for _ in range(6)])\n if load_path:\n if not os.path.exists(load_path):\n import requests\n url = 'https://github.com/JingyunLiang/VRT/releases/download/v0.0/spynet_sintel_final-3d2a1287.pth'\n r = requests.get(url, allow_redirects=True)\n print(f'downloading SpyNet pretrained model from {url}')\n os.makedirs(os.path.dirname(load_path), exist_ok=True)\n open(load_path, 'wb').write(r.content)\n\n self.load_state_dict(torch.load(load_path, map_location=lambda storage, loc: storage)['params'])\n\n self.register_buffer('mean', torch.Tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))\n self.register_buffer('std', torch.Tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))\n\n def preprocess(self, tensor_input):\n tensor_output = (tensor_input - self.mean) / self.std\n return tensor_output\n\n def process(self, ref, supp, w, h, w_floor, h_floor):\n flow_list = []\n\n ref = [self.preprocess(ref)]\n supp = [self.preprocess(supp)]\n\n # ref = [ref]\n # supp = [supp]\n\n for level in range(5):\n ref.insert(0, F.avg_pool2d(input=ref[0], kernel_size=2, stride=2, count_include_pad=False))\n supp.insert(0, F.avg_pool2d(input=supp[0], kernel_size=2, stride=2, count_include_pad=False))\n\n flow = ref[0].new_zeros(\n [ref[0].size(0), 2,\n int(math.floor(ref[0].size(2) / 2.0)),\n int(math.floor(ref[0].size(3) / 2.0))])\n\n for level in range(len(ref)):\n upsampled_flow = F.interpolate(input=flow, scale_factor=2, mode='bilinear', align_corners=True) * 2.0\n\n if upsampled_flow.size(2) != ref[level].size(2):\n upsampled_flow = F.pad(input=upsampled_flow, pad=[0, 0, 0, 1], mode='replicate')\n if upsampled_flow.size(3) != ref[level].size(3):\n upsampled_flow = F.pad(input=upsampled_flow, pad=[0, 1, 0, 0], mode='replicate')\n\n flow = self.basic_module[level](torch.cat([\n ref[level],\n arch_util.flow_warp(\n supp[level], upsampled_flow.permute(0, 2, 3, 1), interp_mode='bilinear', padding_mode='border'),\n upsampled_flow\n ], 1)) + upsampled_flow\n\n if level in self.return_levels:\n scale = 2**(5-level) # level=5 (scale=1), level=4 (scale=2), level=3 (scale=4), level=2 (scale=8)\n flow_out = F.interpolate(input=flow, size=(h//scale, w//scale), mode='bilinear', align_corners=False)\n flow_out[:, 0, :, :] *= float(w//scale) / float(w_floor//scale)\n flow_out[:, 1, :, :] *= float(h//scale) / float(h_floor//scale)\n if torch.abs(flow_out).mean() > 200:\n print(f\"level {level}, flow > 200: {torch.abs(flow_out).mean():.4f}\")\n # return None\n flow_out.clamp(-250, 250)\n flow_list.insert(0, flow_out)\n\n return flow_list\n\n def forward(self, ref, supp):\n assert ref.size() == supp.size()\n\n h, w = ref.size(2), ref.size(3)\n w_floor = math.floor(math.ceil(w / 32.0) * 32.0)\n h_floor = math.floor(math.ceil(h / 32.0) * 32.0)\n\n ref = F.interpolate(input=ref, size=(h_floor, w_floor), mode='bilinear', align_corners=False)\n supp = F.interpolate(input=supp, size=(h_floor, w_floor), mode='bilinear', align_corners=False)\n\n flow_list = self.process(ref, supp, w, h, w_floor, h_floor)\n\n return flow_list[0] if len(flow_list) == 1 else flow_list\n\n\n\nclass FlowGuidedPCDAlign(nn.Module):\n ''' Alignment module using Pyramid, Cascading and Deformable convolution\n with 3 pyramid levels. [From EDVR]\n '''\n\n def __init__(self, nf=64, groups=8):\n super(FlowGuidedPCDAlign, self).__init__()\n # L3: level 3, 1/4 spatial size\n self.L3_offset_conv1 = nn.Conv2d(nf * 2 + 2, nf, 3, 1, 1, bias=True) # concat for diff\n self.L3_offset_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.L3_dcnpack = FlowGuidedDCN(nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups)\n\n # L2: level 2, 1/2 spatial size\n self.L2_offset_conv1 = nn.Conv2d(nf * 2 + 2, nf, 3, 1, 1, bias=True) # concat for diff\n self.L2_offset_conv2 = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True) # concat for offset\n self.L2_offset_conv3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.L2_dcnpack = FlowGuidedDCN(nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups)\n self.L2_fea_conv = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True) # concat for fea\n\n # L1: level 1, original spatial size\n self.L1_offset_conv1 = nn.Conv2d(nf * 2 + 2, nf, 3, 1, 1, bias=True) # concat for diff\n self.L1_offset_conv2 = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True) # concat for offset\n self.L1_offset_conv3 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.L1_dcnpack = FlowGuidedDCN(nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups)\n self.L1_fea_conv = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True) # concat for fea\n\n # Cascading DCN\n self.cas_offset_conv1 = nn.Conv2d(nf * 2, nf, 3, 1, 1, bias=True) # concat for diff\n self.cas_offset_conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)\n self.cas_dcnpack = DCN(nf, nf, 3, stride=1, padding=1, dilation=1, deformable_groups=groups)\n\n self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)\n\n def forward(self, nbr_fea_l, nbr_fea_warped_l, ref_fea_l, flows_l):\n '''align other neighboring frames to the reference frame in the feature level\n nbr_fea_l, ref_fea_l: [L1, L2, L3], each with [B,C,H,W] features\n '''\n # L3\n L3_offset = torch.cat([nbr_fea_warped_l[2], ref_fea_l[2], flows_l[2]], dim=1)\n L3_offset = self.lrelu(self.L3_offset_conv1(L3_offset))\n L3_offset = self.lrelu(self.L3_offset_conv2(L3_offset))\n L3_fea = self.lrelu(self.L3_dcnpack(nbr_fea_l[2], L3_offset, flows_l[2]))\n # L2\n L3_offset = F.interpolate(L3_offset, scale_factor=2, mode='bilinear', align_corners=False)\n L2_offset = torch.cat([nbr_fea_warped_l[1], ref_fea_l[1], flows_l[1]], dim=1)\n L2_offset = self.lrelu(self.L2_offset_conv1(L2_offset))\n L2_offset = self.lrelu(self.L2_offset_conv2(torch.cat([L2_offset, L3_offset*2], dim=1)))\n L2_offset = self.lrelu(self.L2_offset_conv3(L2_offset))\n L2_fea = self.L2_dcnpack(nbr_fea_l[1], L2_offset, flows_l[1])\n L3_fea = F.interpolate(L3_fea, scale_factor=2, mode='bilinear', align_corners=False)\n L2_fea = self.lrelu(self.L2_fea_conv(torch.cat([L2_fea, L3_fea], dim=1)))\n # L1\n L2_offset = F.interpolate(L2_offset, scale_factor=2, mode='bilinear', align_corners=False)\n L1_offset = torch.cat([nbr_fea_warped_l[0], ref_fea_l[0], flows_l[0]], dim=1)\n L1_offset = self.lrelu(self.L1_offset_conv1(L1_offset))\n L1_offset = self.lrelu(self.L1_offset_conv2(torch.cat([L1_offset, L2_offset * 2], dim=1)))\n L1_offset = self.lrelu(self.L1_offset_conv3(L1_offset))\n L1_fea = self.L1_dcnpack(nbr_fea_l[0], L1_offset, flows_l[0])\n L2_fea = F.interpolate(L2_fea, scale_factor=2, mode='bilinear', align_corners=False)\n L1_fea = self.L1_fea_conv(torch.cat([L1_fea, L2_fea], dim=1))\n\n # Cascading\n offset = torch.cat([L1_fea, ref_fea_l[0]], dim=1)\n offset = self.lrelu(self.cas_offset_conv1(offset))\n offset = self.lrelu(self.cas_offset_conv2(offset))\n L1_fea = self.cas_dcnpack(L1_fea, offset)\n\n return L1_fea\n\n\nclass CrossNonLocal_Fusion(nn.Module):\n ''' Cross Non Local fusion module\n '''\n def __init__(self, nf=64, out_feat=96, nframes=5, center=2):\n super(CrossNonLocal_Fusion, self).__init__()\n self.center = center\n\n self.non_local_T = nn.ModuleList()\n self.non_local_F = nn.ModuleList()\n\n for i in range(nframes):\n self.non_local_T.append(NonLocalCross(nf, inter_channels=nf//2, sub_sample=True, bn_layer=False))\n self.non_local_F.append(NonLocal(nf, inter_channels=nf//2, sub_sample=True, bn_layer=False))\n\n # fusion conv: using 1x1 to save parameters and computation\n self.fea_fusion = nn.Conv2d(nframes * nf*2, out_feat, 3, 1, 1, bias=True)\n\n self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)\n\n def forward(self, aligned_fea):\n B, N, C, H, W = aligned_fea.size() # N video frames\n ref = aligned_fea[:, self.center, :, :, :].clone()\n\n cor_l = []\n non_l = []\n for i in range(N):\n nbr = aligned_fea[:, i, :, :, :]\n non_l.append(self.non_local_F[i](nbr))\n cor_l.append(self.non_local_T[i](nbr, ref))\n\n aligned_fea_T = torch.cat(cor_l, dim=1)\n aligned_fea_F = torch.cat(non_l, dim=1)\n aligned_fea = torch.cat([aligned_fea_T, aligned_fea_F], dim=1)\n\n #### fusion\n fea = self.fea_fusion(aligned_fea)\n\n return fea\n\n\n\nclass BSRT(nn.Module):\n def __init__(self, args, nframes=8, img_size=64, patch_size=1, in_chans=3, out_chans=3,\n embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],\n window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,\n drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,\n norm_layer=nn.LayerNorm, ape=False, patch_norm=True,\n use_checkpoint=False, upscale=4, non_local=False,\n **kwargs):\n super(BSRT, self).__init__()\n num_in_ch = in_chans\n num_out_ch = out_chans\n num_feat = 64\n groups = 8\n # embed_dim = num_feat\n back_RBs = 5\n n_resblocks = 6\n\n self.args = args\n self.center = 0\n self.upscale = upscale\n self.window_size = window_size\n self.non_local = non_local\n self.nframes = nframes\n\n self.num_layers = len(depths)\n self.embed_dim = embed_dim\n self.ape = ape\n self.patch_norm = patch_norm\n self.num_features = embed_dim\n self.mlp_ratio = mlp_ratio\n\n spynet_path='/home/luoziwei/.pretrained_models/spynet_sintel_final-3d2a1287.pth'\n self.spynet = SpyNet(spynet_path, [3, 4, 5])\n self.conv_flow = nn.Conv2d(1, 3, kernel_size=3, stride=1, padding=1)\n self.flow_ps = nn.PixelShuffle(2)\n\n # split image into non-overlapping patches\n self.patch_embed = swu.PatchEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,\n norm_layer=norm_layer if self.patch_norm else None)\n num_patches = self.patch_embed.num_patches\n patches_resolution = self.patch_embed.patches_resolution\n self.patches_resolution = patches_resolution\n\n # merge non-overlapping patches into image\n self.patch_unembed = swu.PatchUnEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,\n norm_layer=norm_layer if self.patch_norm else None)\n\n #####################################################################################################\n ################################### 1, shallow feature extraction ###################################\n self.conv_first = nn.Conv2d(num_in_ch*(1+2*0), embed_dim, 3, 1, 1, bias=True)\n \n # # stochastic depth\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule\n\n if args.swinfeature:\n if self.args.local_rank <= 0:\n print(\"using swinfeature\")\n self.pre_layers = nn.ModuleList()\n for i_layer in range(depths[0]):\n layer = swu.SwinTransformerBlock(dim=embed_dim, \n input_resolution=(patches_resolution[0]//2,\n patches_resolution[1]//2),\n num_heads=num_heads[0], window_size=window_size,\n shift_size=0 if (i_layer % 2 == 0) else window_size // 2,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate,\n drop_path=dpr[i_layer],\n norm_layer=norm_layer)\n self.pre_layers.append(layer)\n\n self.pre_norm = norm_layer(embed_dim)\n else:\n WARB = functools.partial(arch_util.WideActResBlock, nf=embed_dim)\n self.feature_extraction = arch_util.make_layer(WARB, 5)\n\n self.conv_after_pre_layer = nn.Conv2d(embed_dim, num_feat*4, 3, 1, 1, bias=True)\n self.mid_ps = nn.PixelShuffle(2)\n\n self.fea_L2_conv1 = nn.Conv2d(num_feat, num_feat*2, 3, 2, 1, bias=True)\n self.fea_L3_conv1 = nn.Conv2d(num_feat*2, num_feat*4, 3, 2, 1, bias=True)\n\n #####################################################################################################\n ################################### 2, Feature Enhanced PCD Align ###################################\n\n # Top layers\n self.toplayer = nn.Conv2d(num_feat*4, num_feat, kernel_size=1, stride=1, padding=0)\n # Smooth layers\n self.smooth1 = nn.Conv2d(num_feat, num_feat, kernel_size=3, stride=1, padding=1)\n self.smooth2 = nn.Conv2d(num_feat, num_feat, kernel_size=3, stride=1, padding=1)\n # Lateral layers\n self.latlayer1 = nn.Conv2d(num_feat*2, num_feat, kernel_size=1, stride=1, padding=0)\n self.latlayer2 = nn.Conv2d(num_feat*1, num_feat, kernel_size=1, stride=1, padding=0)\n\n # self.align = PCD_Align(nf=num_feat, groups=groups)\n self.align = FlowGuidedPCDAlign(nf=num_feat, groups=groups)\n #####################################################################################################\n ################################### 3, Multi-frame Feature Fusion ##################################\n\n if self.non_local:\n if self.args.local_rank <= 0:\n print(\"using non_local\")\n self.fusion = CrossNonLocal_Fusion(nf=num_feat, out_feat=embed_dim, nframes=nframes, center=self.center)\n else:\n self.fusion = nn.Conv2d(nframes * num_feat, embed_dim, 1, 1, bias=True)\n\n #####################################################################################################\n ################################### 4, deep feature extraction ######################################\n\n # absolute position embedding\n if self.ape:\n self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))\n swu.trunc_normal_(self.absolute_pos_embed, std=.02)\n\n self.pos_drop = nn.Dropout(p=drop_rate)\n\n # build Residual Swin Transformer blocks (RSTB)\n self.layers = nn.ModuleList()\n for i_layer in range(1, self.num_layers):\n layer = swu.RSTB(dim=embed_dim,\n input_resolution=(patches_resolution[0],\n patches_resolution[1]),\n depth=depths[i_layer],\n num_heads=num_heads[i_layer],\n window_size=window_size,\n mlp_ratio=self.mlp_ratio,\n qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate,\n drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results\n norm_layer=norm_layer,\n downsample=None,\n use_checkpoint=use_checkpoint,\n img_size=img_size,\n patch_size=patch_size\n )\n self.layers.append(layer)\n \n self.norm = norm_layer(self.num_features)\n\n # build the last conv layer in deep feature extraction\n self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)\n\n #####################################################################################################\n ################################ 5, high quality image reconstruction ################################\n\n self.upconv1 = nn.Conv2d(embed_dim, num_feat * 4, 3, 1, 1, bias=True)\n self.upconv2 = nn.Conv2d(num_feat, 64 * 4, 3, 1, 1, bias=True)\n self.pixel_shuffle = nn.PixelShuffle(2)\n self.HRconv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)\n self.conv_last = nn.Conv2d(64, args.n_colors, 3, 1, 1, bias=True)\n\n #### skip #############\n self.skip_pixel_shuffle = nn.PixelShuffle(2)\n self.skipup1 = nn.Conv2d(num_in_ch//4, num_feat * 4, 3, 1, 1, bias=True)\n self.skipup2 = nn.Conv2d(num_feat, args.n_colors * 4, 3, 1, 1, bias=True)\n\n #### activation function\n self.lrelu = nn.LeakyReLU(0.1, inplace=True)\n self.lrelu2 = nn.LeakyReLU(0.1, inplace=True)\n\n self.apply(self._init_weights)\n\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n swu.trunc_normal_(m.weight, std=.02)\n if isinstance(m, nn.Linear) and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n nn.init.constant_(m.bias, 0)\n nn.init.constant_(m.weight, 1.0)\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'absolute_pos_embed'}\n\n @torch.jit.ignore\n def no_weight_decay_keywords(self):\n return {'relative_position_bias_table'}\n\n def _upsample_add(self, x, y):\n return F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False) + y\n\n def check_image_size(self, x):\n _, _, h, w = x.size()\n mod_pad_h = (self.window_size - h % self.window_size) % self.window_size\n mod_pad_w = (self.window_size - w % self.window_size) % self.window_size\n x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')\n return x\n\n def pre_forward_features(self, x):\n if self.args.swinfeature:\n x_size = (x.shape[-2], x.shape[-1])\n x = self.patch_embed(x, use_norm=True)\n if self.ape:\n x = x + self.absolute_pos_embed\n x = self.pos_drop(x)\n\n for idx, layer in enumerate(self.pre_layers):\n x = layer(x, x_size)\n\n x = self.pre_norm(x)\n x = self.patch_unembed(x, x_size)\n\n else:\n x = self.feature_extraction(x)\n\n return x\n\n def forward_features(self, x):\n x_size = (x.shape[-2], x.shape[-1])\n x = self.patch_embed(x)\n if self.ape:\n x = x + self.absolute_pos_embed\n x = self.pos_drop(x)\n\n for idx, layer in enumerate(self.layers):\n x = layer(x, x_size)\n if torch.any(torch.isinf(x)) or torch.any(torch.isnan(x)):\n print('layer: ', idx)\n\n x = self.norm(x) # B L C\n x = self.patch_unembed(x, x_size)\n\n return x\n\n @autocast()\n def forward(self, x, print_time=False):\n B, N, C, H, W = x.size() # N video frames\n x_center = x[:, self.center, :, :, :].contiguous()\n\n #### skip module ########\n skip1 = self.lrelu2(self.skip_pixel_shuffle(self.skipup1(self.skip_pixel_shuffle(x_center))))\n skip2 = self.skip_pixel_shuffle(self.skipup2(skip1))\n\n x_ = self.conv_flow(self.flow_ps(x.view(B*N, C, H, W))).view(B, N, -1, H*2, W*2)\n \n # calculate flows\n ref_flows = self.get_ref_flows(x_)\n\n #### extract LR features\n x = self.lrelu(self.conv_first(x.view(B*N, -1, H, W)))\n\n L1_fea = self.mid_ps(self.conv_after_pre_layer(self.pre_forward_features(x)))\n _, _, H, W = L1_fea.size()\n\n L2_fea = self.lrelu(self.fea_L2_conv1(L1_fea))\n L3_fea = self.lrelu(self.fea_L3_conv1(L2_fea))\n\n # FPN enhance features\n L3_fea = self.lrelu(self.toplayer(L3_fea))\n L2_fea = self.smooth1(self._upsample_add(L3_fea, self.latlayer1(L2_fea)))\n L1_fea = self.smooth2(self._upsample_add(L2_fea, self.latlayer2(L1_fea)))\n\n L1_fea = L1_fea.view(B, N, -1, H, W).contiguous()\n L2_fea = L2_fea.view(B, N, -1, H // 2, W // 2 ).contiguous()\n L3_fea = L3_fea.view(B, N, -1, H // 4, W // 4).contiguous()\n\n #### PCD align\n # ref feature list\n ref_fea_l = [\n L1_fea[:, self.center, :, :, :].clone(), \n L2_fea[:, self.center, :, :, :].clone(),\n L3_fea[:, self.center, :, :, :].clone()\n ]\n aligned_fea = []\n for i in range(N):\n nbr_fea_l = [\n L1_fea[:, i, :, :, :].clone(), \n L2_fea[:, i, :, :, :].clone(),\n L3_fea[:, i, :, :, :].clone()\n ]\n flows_l = [\n ref_flows[0][:, i, :, :, :].clone(), \n ref_flows[1][:, i, :, :, :].clone(), \n ref_flows[2][:, i, :, :, :].clone()\n ]\n # print(nbr_fea_l[0].shape, flows_l[0].shape)\n nbr_warped_l = [\n arch_util.flow_warp(nbr_fea_l[0], flows_l[0].permute(0, 2, 3, 1), 'bilinear'),\n arch_util.flow_warp(nbr_fea_l[1], flows_l[1].permute(0, 2, 3, 1), 'bilinear'),\n arch_util.flow_warp(nbr_fea_l[2], flows_l[2].permute(0, 2, 3, 1), 'bilinear')\n ]\n aligned_fea.append(self.align(nbr_fea_l, nbr_warped_l, ref_fea_l, flows_l))\n\n aligned_fea = torch.stack(aligned_fea, dim=1) # [B, N, C, H, W] --> [B, T, C, H, W]\n\n if not self.non_local:\n aligned_fea = aligned_fea.view(B, -1, H, W)\n\n x = self.lrelu(self.fusion(aligned_fea))\n\n x = self.lrelu(self.conv_after_body(self.forward_features(x))) + x\n\n x = self.lrelu(self.pixel_shuffle(self.upconv1(x)))\n x = skip1 + x\n x = self.lrelu(self.pixel_shuffle(self.upconv2(x)))\n x = self.lrelu(self.HRconv(x))\n x = self.conv_last(x)\n\n x = skip2 + x\n return x\n\n\n def get_ref_flows(self, x):\n '''Get flow between frames ref and other'''\n\n b, n, c, h, w = x.size()\n x_nbr = x.reshape(-1, c, h, w)\n x_ref = x[:, self.center:self.center+1, :, :, :].repeat(1, n, 1, 1, 1).reshape(-1, c, h, w)\n\n # backward\n flows = self.spynet(x_ref, x_nbr)\n flows_list = [flow.view(b, n, 2, h // (2 ** (i)), w // (2 ** (i))) for flow, i in\n zip(flows, range(3))]\n\n return flows_list\n\n\n\n\n\n\n\n"
},
{
"path": "code/real/bsrt/model/checkpoint.py",
"content": "import torch\nimport warnings\n\n\ndef detach_variable(inputs):\n if isinstance(inputs, tuple):\n out = []\n for inp in inputs:\n x = inp.detach()\n x.requires_grad = inp.requires_grad\n out.append(x)\n return tuple(out)\n else:\n raise RuntimeError(\n \"Only tuple of tensors is supported. Got Unsupported input type: \", type(inputs).__name__)\n\n\ndef check_backward_validity(inputs):\n if not any(inp.requires_grad for inp in inputs):\n warnings.warn(\"None of the inputs have requires_grad=True. Gradients will be None\")\n\n\nclass CheckpointFunction(torch.autograd.Function):\n @staticmethod\n def forward(ctx, run_function, length, *args):\n ctx.run_function = run_function\n ctx.input_tensors = list(args[:length])\n ctx.input_params = list(args[length:])\n with torch.no_grad():\n output_tensors = ctx.run_function(*ctx.input_tensors)\n return output_tensors\n\n @staticmethod\n def backward(ctx, *output_grads):\n for i in range(len(ctx.input_tensors)):\n temp = ctx.input_tensors[i]\n ctx.input_tensors[i] = temp.detach()\n ctx.input_tensors[i].requires_grad = temp.requires_grad\n with torch.enable_grad():\n output_tensors = ctx.run_function(*ctx.input_tensors)\n input_grads = torch.autograd.grad(output_tensors, ctx.input_tensors + ctx.input_params, output_grads, allow_unused=True)\n return (None, None) + input_grads\n"
},
{
"path": "code/real/bsrt/model/common.py",
"content": "import math\nimport numpy as np\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n\ndef default_conv(in_channels, out_channels, kernel_size, bias=True):\n return nn.Conv2d(\n in_channels, out_channels, kernel_size,\n padding=(kernel_size // 2), bias=bias)\n\n\nclass MeanShift(nn.Conv2d):\n def __init__(\n self, rgb_range,\n rgb_mean=(0.4488, 0.4371, 0.4040), rgb_std=(1.0, 1.0, 1.0), sign=-1):\n super(MeanShift, self).__init__(3, 3, kernel_size=1)\n std = torch.Tensor(rgb_std)\n self.weight.data = torch.eye(3).view(3, 3, 1, 1) / std.view(3, 1, 1, 1)\n self.bias.data = sign * rgb_range * torch.Tensor(rgb_mean) / std\n for p in self.parameters():\n p.requires_grad = False\n\n\nclass BasicBlock(nn.Sequential):\n def __init__(\n self, conv, in_channels, out_channels, kernel_size, stride=1, bias=False,\n bn=True, act=nn.ReLU(True)):\n\n m = [conv(in_channels, out_channels, kernel_size, bias=bias)]\n if bn:\n m.append(nn.BatchNorm2d(out_channels))\n if act is not None:\n m.append(act)\n\n super(BasicBlock, self).__init__(*m)\n\n\nclass ResBlock(nn.Module):\n def __init__(\n self, conv, n_feats, kernel_size,\n bias=True, bn=False, act=nn.ReLU(True), res_scale=1):\n\n super(ResBlock, self).__init__()\n m = []\n for i in range(2):\n m.append(conv(n_feats, n_feats, kernel_size, bias=bias))\n if bn:\n m.append(nn.BatchNorm2d(n_feats))\n if i == 0:\n m.append(act)\n\n self.body = nn.Sequential(*m)\n self.res_scale = res_scale\n\n def forward(self, x):\n res = self.body(x).mul(self.res_scale)\n res += x\n\n return res\n\n\nclass Upsampler(nn.Sequential):\n def __init__(self, conv, scale, n_feats, bn=False, act=False, bias=True):\n\n m = []\n if (scale & (scale - 1)) == 0: # Is scale = 2^n?\n for _ in range(int(math.log(scale, 2))):\n m.append(conv(n_feats, 4 * n_feats, 3, bias))\n m.append(nn.PixelShuffle(2))\n if bn:\n m.append(nn.BatchNorm2d(n_feats))\n if act == 'relu':\n m.append(nn.ReLU(True))\n elif act == 'prelu':\n m.append(nn.PReLU(n_feats))\n\n elif scale == 3:\n m.append(conv(n_feats, 9 * n_feats, 3, bias))\n m.append(nn.PixelShuffle(3))\n if bn:\n m.append(nn.BatchNorm2d(n_feats))\n if act == 'relu':\n m.append(nn.ReLU(True))\n elif act == 'prelu':\n m.append(nn.PReLU(n_feats))\n else:\n raise NotImplementedError\n\n super(Upsampler, self).__init__(*m)\n\n\nclass UpOnly(nn.Sequential):\n def __init__(self, scale):\n\n m = []\n if (scale & (scale - 1)) == 0: # Is scale = 2^n?\n for _ in range(int(math.log(scale, 2))):\n m.append(nn.PixelShuffle(2))\n\n\n elif scale == 3:\n\n m.append(nn.PixelShuffle(3))\n\n else:\n raise NotImplementedError\n\n super(UpOnly, self).__init__(*m)\n\n\ndef lanczos_kernel(dx, a=3, N=None, dtype=None, device=None):\n '''\n Generates 1D Lanczos kernels for translation and interpolation.\n Args:\n dx : float, tensor (batch_size, 1), the translation in pixels to shift an image.\n a : int, number of lobes in the kernel support.\n If N is None, then the width is the kernel support (length of all lobes),\n S = 2(a + ceil(dx)) + 1.\n N : int, width of the kernel.\n If smaller than S then N is set to S.\n Returns:\n k: tensor (?, ?), lanczos kernel\n '''\n\n if not torch.is_tensor(dx):\n dx = torch.tensor(dx, dtype=dtype, device=device)\n\n if device is None:\n device = dx.device\n\n if dtype is None:\n dtype = dx.dtype\n\n D = dx.abs().ceil().int()\n S = 2 * (a + D) + 1 # width of kernel support\n\n S_max = S.max() if hasattr(S, 'shape') else S\n\n if (N is None) or (N < S_max):\n N = S\n\n Z = (N - S) // 2 # width of zeros beyond kernel support\n\n start = (-(a + D + Z)).min()\n end = (a + D + Z + 1).max()\n x = torch.arange(start, end, dtype=dtype, device=device).view(1, -1) - dx\n px = (np.pi * x) + 1e-3\n\n sin_px = torch.sin(px)\n sin_pxa = torch.sin(px / a)\n\n k = a * sin_px * sin_pxa / px ** 2 # sinc(x) masked by sinc(x/a)\n\n return k\n\n\ndef lanczos_shift(img, shift, p=5, a=3):\n '''\n Shifts an image by convolving it with a Lanczos kernel.\n Lanczos interpolation is an approximation to ideal sinc interpolation,\n by windowing a sinc kernel with another sinc function extending up to a\n few nunber of its lobes (typically a=3).\n\n Args:\n img : tensor (batch_size, channels, height, width), the images to be shifted\n shift : tensor (batch_size, 2) of translation parameters (dy, dx)\n p : int, padding width prior to convolution (default=3)\n a : int, number of lobes in the Lanczos interpolation kernel (default=3)\n Returns:\n I_s: tensor (batch_size, channels, height, width), shifted images\n '''\n img = img.transpose(0, 1)\n dtype = img.dtype\n\n if len(img.shape) == 2:\n img = img[None, None].repeat(1, shift.shape[0], 1, 1) # batch of one image\n elif len(img.shape) == 3: # one image per shift\n assert img.shape[0] == shift.shape[0]\n img = img[None,]\n\n # Apply padding\n\n padder = torch.nn.ReflectionPad2d(p) # reflect pre-padding\n I_padded = padder(img)\n\n # Create 1D shifting kernels\n\n y_shift = shift[:, [0]]\n x_shift = shift[:, [1]]\n\n k_y = (lanczos_kernel(y_shift, a=a, N=None, dtype=dtype)\n .flip(1) # flip axis of convolution\n )[:, None, :, None] # expand dims to get shape (batch, channels, y_kernel, 1)\n k_x = (lanczos_kernel(x_shift, a=a, N=None, dtype=dtype)\n .flip(1)\n )[:, None, None, :] # shape (batch, channels, 1, x_kernel)\n\n # Apply kernels\n # print(I_padded.shape, k_y.shape)\n I_s = torch.conv1d(I_padded,\n groups=k_y.shape[0],\n weight=k_y,\n padding=[k_y.shape[2] // 2, 0]) # same padding\n I_s = torch.conv1d(I_s,\n groups=k_x.shape[0],\n weight=k_x,\n padding=[0, k_x.shape[3] // 2])\n\n I_s = I_s[..., p:-p, p:-p] # remove padding\n\n # print(I_s.shape)\n return I_s.transpose(0, 1) # , k.squeeze()\n"
},
{
"path": "code/real/bsrt/model/non_local/network.py",
"content": "from torch import nn\n# from lib.non_local_concatenation import NONLocalBlock2D\n# from lib.non_local_gaussian import NONLocalBlock2D\nfrom lib.non_local_embedded_gaussian import NONLocalBlock2D\n# from lib.non_local_dot_product import NONLocalBlock2D\n\n\nclass Network(nn.Module):\n def __init__(self):\n super(Network, self).__init__()\n\n self.conv_1 = nn.Sequential(\n nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(32),\n nn.ReLU(),\n nn.MaxPool2d(2),\n )\n\n self.nl_1 = NONLocalBlock2D(in_channels=32)\n self.conv_2 = nn.Sequential(\n nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(64),\n nn.ReLU(),\n nn.MaxPool2d(2),\n )\n\n self.nl_2 = NONLocalBlock2D(in_channels=64)\n self.conv_3 = nn.Sequential(\n nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, stride=1, padding=1),\n nn.BatchNorm2d(128),\n nn.ReLU(),\n nn.MaxPool2d(2),\n )\n\n self.fc = nn.Sequential(\n nn.Linear(in_features=128*3*3, out_features=256),\n nn.ReLU(),\n nn.Dropout(0.5),\n\n nn.Linear(in_features=256, out_features=10)\n )\n\n def forward(self, x):\n batch_size = x.size(0)\n\n feature_1 = self.conv_1(x)\n nl_feature_1 = self.nl_1(feature_1)\n\n feature_2 = self.conv_2(nl_feature_1)\n nl_feature_2 = self.nl_2(feature_2)\n\n output = self.conv_3(nl_feature_2).view(batch_size, -1)\n output = self.fc(output)\n\n return output\n\n def forward_with_nl_map(self, x):\n batch_size = x.size(0)\n\n feature_1 = self.conv_1(x)\n nl_feature_1, nl_map_1 = self.nl_1(feature_1, return_nl_map=True)\n\n feature_2 = self.conv_2(nl_feature_1)\n nl_feature_2, nl_map_2 = self.nl_2(feature_2, return_nl_map=True)\n\n output = self.conv_3(nl_feature_2).view(batch_size, -1)\n output = self.fc(output)\n\n return output, [nl_map_1, nl_map_2]\n\n\nif __name__ == '__main__':\n import torch\n\n img = torch.randn(3, 1, 28, 28)\n net = Network()\n out = net(img)\n print(out.size())\n\n"
},
{
"path": "code/real/bsrt/model/non_local/non_local_concatenation.py",
"content": "import torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\nclass _NonLocalBlockND(nn.Module):\n def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):\n super(_NonLocalBlockND, self).__init__()\n\n assert dimension in [1, 2, 3]\n\n self.dimension = dimension\n self.sub_sample = sub_sample\n\n self.in_channels = in_channels\n self.inter_channels = inter_channels\n\n if self.inter_channels is None:\n self.inter_channels = in_channels // 2\n if self.inter_channels == 0:\n self.inter_channels = 1\n\n if dimension == 3:\n conv_nd = nn.Conv3d\n max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))\n bn = nn.BatchNorm3d\n elif dimension == 2:\n conv_nd = nn.Conv2d\n max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))\n bn = nn.BatchNorm2d\n else:\n conv_nd = nn.Conv1d\n max_pool_layer = nn.MaxPool1d(kernel_size=(2))\n bn = nn.BatchNorm1d\n\n self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if bn_layer:\n self.W = nn.Sequential(\n conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0),\n bn(self.in_channels)\n )\n nn.init.constant_(self.W[1].weight, 0)\n nn.init.constant_(self.W[1].bias, 0)\n else:\n self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0)\n nn.init.constant_(self.W.weight, 0)\n nn.init.constant_(self.W.bias, 0)\n\n self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n self.concat_project = nn.Sequential(\n nn.Conv2d(self.inter_channels * 2, 1, 1, 1, 0, bias=False),\n nn.ReLU()\n )\n\n if sub_sample:\n self.g = nn.Sequential(self.g, max_pool_layer)\n self.phi = nn.Sequential(self.phi, max_pool_layer)\n\n def forward(self, x, return_nl_map=False):\n '''\n :param x: (b, c, t, h, w)\n :param return_nl_map: if True return z, nl_map, else only return z.\n :return:\n '''\n\n batch_size = x.size(0)\n\n g_x = self.g(x).view(batch_size, self.inter_channels, -1)\n g_x = g_x.permute(0, 2, 1)\n\n # (b, c, N, 1)\n theta_x = self.theta(x).view(batch_size, self.inter_channels, -1, 1)\n # (b, c, 1, N)\n phi_x = self.phi(x).view(batch_size, self.inter_channels, 1, -1)\n\n h = theta_x.size(2)\n w = phi_x.size(3)\n theta_x = theta_x.repeat(1, 1, 1, w)\n phi_x = phi_x.repeat(1, 1, h, 1)\n\n concat_feature = torch.cat([theta_x, phi_x], dim=1)\n f = self.concat_project(concat_feature)\n b, _, h, w = f.size()\n f = f.view(b, h, w)\n\n N = f.size(-1)\n f_div_C = f / N\n\n y = torch.matmul(f_div_C, g_x)\n y = y.permute(0, 2, 1).contiguous()\n y = y.view(batch_size, self.inter_channels, *x.size()[2:])\n W_y = self.W(y)\n z = W_y + x\n\n if return_nl_map:\n return z, f_div_C\n return z\n\n\nclass NONLocalBlock1D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock1D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=1, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock2D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock2D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=2, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock3D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True,):\n super(NONLocalBlock3D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=3, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nif __name__ == '__main__':\n import torch\n\n for (sub_sample_, bn_layer_) in [(True, True), (False, False), (True, False), (False, True)]:\n img = torch.zeros(2, 3, 20)\n net = NONLocalBlock1D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.zeros(2, 3, 20, 20)\n net = NONLocalBlock2D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.randn(2, 3, 8, 20, 20)\n net = NONLocalBlock3D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n"
},
{
"path": "code/real/bsrt/model/non_local/non_local_cross_dot_product.py",
"content": "import torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\nclass _NonLocalBlockND(nn.Module):\n def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):\n super(_NonLocalBlockND, self).__init__()\n\n assert dimension in [1, 2, 3]\n\n self.dimension = dimension\n self.sub_sample = sub_sample\n\n self.in_channels = in_channels\n self.inter_channels = inter_channels\n\n if self.inter_channels is None:\n self.inter_channels = in_channels // 2\n if self.inter_channels == 0:\n self.inter_channels = 1\n\n if dimension == 3:\n conv_nd = nn.Conv3d\n max_pool_layer = nn.MaxPool3d(kernel_size=(1, 4, 4))\n bn = nn.BatchNorm3d\n elif dimension == 2:\n conv_nd = nn.Conv2d\n max_pool_layer = nn.MaxPool2d(kernel_size=(4, 4))\n bn = nn.BatchNorm2d\n else:\n conv_nd = nn.Conv1d\n max_pool_layer = nn.MaxPool1d(kernel_size=(4))\n bn = nn.BatchNorm1d\n\n self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if bn_layer:\n self.W = nn.Sequential(\n conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0),\n bn(self.in_channels)\n )\n nn.init.constant_(self.W[1].weight, 0)\n nn.init.constant_(self.W[1].bias, 0)\n else:\n self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0)\n nn.init.constant_(self.W.weight, 0)\n nn.init.constant_(self.W.bias, 0)\n\n self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if sub_sample:\n self.g = nn.Sequential(self.g, max_pool_layer)\n self.phi = nn.Sequential(self.phi, max_pool_layer)\n\n def forward(self, x, ref, return_nl_map=False):\n \"\"\"\n :param x: (b, c, t, h, w)\n :param return_nl_map: if True return z, nl_map, else only return z.\n :return:\n \"\"\"\n\n batch_size = x.size(0)\n\n g_x = self.g(x).view(batch_size, self.inter_channels, -1)\n g_x = g_x.permute(0, 2, 1)\n\n theta_ref = self.theta(ref).view(batch_size, self.inter_channels, -1)\n theta_ref = theta_ref.permute(0, 2, 1)\n phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)\n f = torch.matmul(theta_ref, phi_x)\n N = f.size(-1)\n f_div_C = f / N\n\n y = torch.matmul(f_div_C, g_x)\n y = y.permute(0, 2, 1).contiguous()\n y = y.view(batch_size, self.inter_channels, *x.size()[2:])\n W_y = self.W(y)\n z = W_y + x\n\n if return_nl_map:\n return z, f_div_C\n return z\n\n\nclass NONLocalBlock1D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock1D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=1, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock2D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock2D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=2, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock3D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock3D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=3, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nif __name__ == '__main__':\n import torch\n\n for (sub_sample_, bn_layer_) in [(True, True), (False, False), (True, False), (False, True)]:\n img = torch.zeros(2, 3, 20)\n net = NONLocalBlock1D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.zeros(2, 3, 20, 20)\n net = NONLocalBlock2D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.randn(2, 3, 8, 20, 20)\n net = NONLocalBlock3D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n\n\n"
},
{
"path": "code/real/bsrt/model/non_local/non_local_dot_product.py",
"content": "import torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\nclass _NonLocalBlockND(nn.Module):\n def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):\n super(_NonLocalBlockND, self).__init__()\n\n assert dimension in [1, 2, 3]\n\n self.dimension = dimension\n self.sub_sample = sub_sample\n\n self.in_channels = in_channels\n self.inter_channels = inter_channels\n\n if self.inter_channels is None:\n self.inter_channels = in_channels // 2\n if self.inter_channels == 0:\n self.inter_channels = 1\n\n if dimension == 3:\n conv_nd = nn.Conv3d\n max_pool_layer = nn.MaxPool3d(kernel_size=(1, 4, 4))\n bn = nn.BatchNorm3d\n elif dimension == 2:\n conv_nd = nn.Conv2d\n max_pool_layer = nn.MaxPool2d(kernel_size=(4, 4))\n bn = nn.BatchNorm2d\n else:\n conv_nd = nn.Conv1d\n max_pool_layer = nn.MaxPool1d(kernel_size=(2))\n bn = nn.BatchNorm1d\n\n self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if bn_layer:\n self.W = nn.Sequential(\n conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0),\n bn(self.in_channels)\n )\n nn.init.constant_(self.W[1].weight, 0)\n nn.init.constant_(self.W[1].bias, 0)\n else:\n self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0)\n nn.init.constant_(self.W.weight, 0)\n nn.init.constant_(self.W.bias, 0)\n\n self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if sub_sample:\n self.g = nn.Sequential(self.g, max_pool_layer)\n self.phi = nn.Sequential(self.phi, max_pool_layer)\n\n def forward(self, x, return_nl_map=False):\n \"\"\"\n :param x: (b, c, t, h, w)\n :param return_nl_map: if True return z, nl_map, else only return z.\n :return:\n \"\"\"\n\n batch_size = x.size(0)\n\n g_x = self.g(x).view(batch_size, self.inter_channels, -1)\n g_x = g_x.permute(0, 2, 1)\n\n theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)\n theta_x = theta_x.permute(0, 2, 1)\n phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)\n f = torch.matmul(theta_x, phi_x)\n N = f.size(-1)\n f_div_C = f / N\n\n y = torch.matmul(f_div_C, g_x)\n y = y.permute(0, 2, 1).contiguous()\n y = y.view(batch_size, self.inter_channels, *x.size()[2:])\n W_y = self.W(y)\n z = W_y + x\n\n if return_nl_map:\n return z, f_div_C\n return z\n\n\nclass NONLocalBlock1D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock1D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=1, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock2D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock2D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=2, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock3D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock3D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=3, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nif __name__ == '__main__':\n import torch\n\n for (sub_sample_, bn_layer_) in [(True, True), (False, False), (True, False), (False, True)]:\n img = torch.zeros(2, 3, 20)\n net = NONLocalBlock1D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.zeros(2, 3, 20, 20)\n net = NONLocalBlock2D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.randn(2, 3, 8, 20, 20)\n net = NONLocalBlock3D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n\n\n"
},
{
"path": "code/real/bsrt/model/non_local/non_local_embedded_gaussian.py",
"content": "import torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\nclass _NonLocalBlockND(nn.Module):\n def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):\n \"\"\"\n :param in_channels:\n :param inter_channels:\n :param dimension:\n :param sub_sample:\n :param bn_layer:\n \"\"\"\n\n super(_NonLocalBlockND, self).__init__()\n\n assert dimension in [1, 2, 3]\n\n self.dimension = dimension\n self.sub_sample = sub_sample\n\n self.in_channels = in_channels\n self.inter_channels = inter_channels\n\n if self.inter_channels is None:\n self.inter_channels = in_channels // 2\n if self.inter_channels == 0:\n self.inter_channels = 1\n\n if dimension == 3:\n conv_nd = nn.Conv3d\n max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))\n bn = nn.BatchNorm3d\n elif dimension == 2:\n conv_nd = nn.Conv2d\n max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))\n bn = nn.BatchNorm2d\n else:\n conv_nd = nn.Conv1d\n max_pool_layer = nn.MaxPool1d(kernel_size=(2))\n bn = nn.BatchNorm1d\n\n self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if bn_layer:\n self.W = nn.Sequential(\n conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0),\n bn(self.in_channels)\n )\n nn.init.constant_(self.W[1].weight, 0)\n nn.init.constant_(self.W[1].bias, 0)\n else:\n self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0)\n nn.init.constant_(self.W.weight, 0)\n nn.init.constant_(self.W.bias, 0)\n\n self.theta = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n self.phi = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if sub_sample:\n self.g = nn.Sequential(self.g, max_pool_layer)\n self.phi = nn.Sequential(self.phi, max_pool_layer)\n\n def forward(self, x, return_nl_map=False):\n \"\"\"\n :param x: (b, c, t, h, w)\n :param return_nl_map: if True return z, nl_map, else only return z.\n :return:\n \"\"\"\n\n batch_size = x.size(0)\n\n g_x = self.g(x).view(batch_size, self.inter_channels, -1)\n g_x = g_x.permute(0, 2, 1)\n\n theta_x = self.theta(x).view(batch_size, self.inter_channels, -1)\n theta_x = theta_x.permute(0, 2, 1)\n phi_x = self.phi(x).view(batch_size, self.inter_channels, -1)\n f = torch.matmul(theta_x, phi_x)\n f_div_C = F.softmax(f, dim=-1)\n\n y = torch.matmul(f_div_C, g_x)\n y = y.permute(0, 2, 1).contiguous()\n y = y.view(batch_size, self.inter_channels, *x.size()[2:])\n W_y = self.W(y)\n z = W_y + x\n\n if return_nl_map:\n return z, f_div_C\n return z\n\n\nclass NONLocalBlock1D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock1D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=1, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock2D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock2D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=2, sub_sample=sub_sample,\n bn_layer=bn_layer,)\n\n\nclass NONLocalBlock3D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock3D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=3, sub_sample=sub_sample,\n bn_layer=bn_layer,)\n\n\nif __name__ == '__main__':\n import torch\n\n for (sub_sample_, bn_layer_) in [(True, True), (False, False), (True, False), (False, True)]:\n img = torch.zeros(2, 3, 20)\n net = NONLocalBlock1D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.zeros(2, 3, 20, 20)\n net = NONLocalBlock2D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.randn(2, 3, 8, 20, 20)\n net = NONLocalBlock3D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n\n"
},
{
"path": "code/real/bsrt/model/non_local/non_local_gaussian.py",
"content": "import torch\nfrom torch import nn\nfrom torch.nn import functional as F\n\n\nclass _NonLocalBlockND(nn.Module):\n def __init__(self, in_channels, inter_channels=None, dimension=3, sub_sample=True, bn_layer=True):\n super(_NonLocalBlockND, self).__init__()\n\n assert dimension in [1, 2, 3]\n\n self.dimension = dimension\n self.sub_sample = sub_sample\n\n self.in_channels = in_channels\n self.inter_channels = inter_channels\n\n if self.inter_channels is None:\n self.inter_channels = in_channels // 2\n if self.inter_channels == 0:\n self.inter_channels = 1\n\n if dimension == 3:\n conv_nd = nn.Conv3d\n max_pool_layer = nn.MaxPool3d(kernel_size=(1, 2, 2))\n bn = nn.BatchNorm3d\n elif dimension == 2:\n conv_nd = nn.Conv2d\n max_pool_layer = nn.MaxPool2d(kernel_size=(2, 2))\n bn = nn.BatchNorm2d\n else:\n conv_nd = nn.Conv1d\n max_pool_layer = nn.MaxPool1d(kernel_size=(2))\n bn = nn.BatchNorm1d\n\n self.g = conv_nd(in_channels=self.in_channels, out_channels=self.inter_channels,\n kernel_size=1, stride=1, padding=0)\n\n if bn_layer:\n self.W = nn.Sequential(\n conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0),\n bn(self.in_channels)\n )\n nn.init.constant_(self.W[1].weight, 0)\n nn.init.constant_(self.W[1].bias, 0)\n else:\n self.W = conv_nd(in_channels=self.inter_channels, out_channels=self.in_channels,\n kernel_size=1, stride=1, padding=0)\n nn.init.constant_(self.W.weight, 0)\n nn.init.constant_(self.W.bias, 0)\n\n if sub_sample:\n self.g = nn.Sequential(self.g, max_pool_layer)\n self.phi = max_pool_layer\n\n def forward(self, x, return_nl_map=False):\n \"\"\"\n :param x: (b, c, t, h, w)\n :param return_nl_map: if True return z, nl_map, else only return z.\n :return:\n \"\"\"\n\n batch_size = x.size(0)\n\n g_x = self.g(x).view(batch_size, self.inter_channels, -1)\n\n g_x = g_x.permute(0, 2, 1)\n\n theta_x = x.view(batch_size, self.in_channels, -1)\n theta_x = theta_x.permute(0, 2, 1)\n\n if self.sub_sample:\n phi_x = self.phi(x).view(batch_size, self.in_channels, -1)\n else:\n phi_x = x.view(batch_size, self.in_channels, -1)\n\n f = torch.matmul(theta_x, phi_x)\n f_div_C = F.softmax(f, dim=-1)\n\n # if self.store_last_batch_nl_map:\n # self.nl_map = f_div_C\n\n y = torch.matmul(f_div_C, g_x)\n y = y.permute(0, 2, 1).contiguous()\n y = y.view(batch_size, self.inter_channels, *x.size()[2:])\n W_y = self.W(y)\n z = W_y + x\n\n if return_nl_map:\n return z, f_div_C\n return z\n\n\nclass NONLocalBlock1D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock1D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=1, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock2D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock2D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=2, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nclass NONLocalBlock3D(_NonLocalBlockND):\n def __init__(self, in_channels, inter_channels=None, sub_sample=True, bn_layer=True):\n super(NONLocalBlock3D, self).__init__(in_channels,\n inter_channels=inter_channels,\n dimension=3, sub_sample=sub_sample,\n bn_layer=bn_layer)\n\n\nif __name__ == '__main__':\n import torch\n\n for (sub_sample_, bn_layer_) in [(True, True), (False, False), (True, False), (False, True)]:\n img = torch.zeros(2, 3, 20)\n net = NONLocalBlock1D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.zeros(2, 3, 20, 20)\n net = NONLocalBlock2D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n img = torch.randn(2, 3, 8, 20, 20)\n net = NONLocalBlock3D(3, sub_sample=sub_sample_, bn_layer=bn_layer_)\n out = net(img)\n print(out.size())\n\n\n\n\n\n\n"
},
{
"path": "code/real/bsrt/model/swin_util.py",
"content": "# -----------------------------------------------------------------------------------\n# SwinIR: Image Restoration Using Swin Transformer, https://arxiv.org/abs/2108.10257\n# Originally Written by Ze Liu, Modified by Jingyun Liang.\n# -----------------------------------------------------------------------------------\n\nimport math\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n# import torch.utils.checkpoint as checkpoint\nfrom model.checkpoint import CheckpointFunction as checkpoint\nfrom timm.models.layers import DropPath, to_2tuple, trunc_normal_\nimport time\nfrom functools import reduce, lru_cache\n\n\nclass Mlp(nn.Module):\n def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n self.fc1 = nn.Linear(in_features, hidden_features)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x):\n x = self.fc1(x)\n x = self.act(x)\n x = self.drop(x)\n x = self.fc2(x)\n x = self.drop(x)\n return x\n\nclass Mlp_GEGLU(nn.Module):\n \"\"\" Multilayer perceptron with gated linear unit (GEGLU). Ref. \"GLU Variants Improve Transformer\".\n\n Args:\n x: (B, D, H, W, C)\n\n Returns:\n x: (B, D, H, W, C)\n \"\"\"\n\n def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):\n super().__init__()\n out_features = out_features or in_features\n hidden_features = hidden_features or in_features\n\n self.fc11 = nn.Linear(in_features, hidden_features)\n self.fc12 = nn.Linear(in_features, hidden_features)\n self.act = act_layer()\n self.fc2 = nn.Linear(hidden_features, out_features)\n self.drop = nn.Dropout(drop)\n\n def forward(self, x):\n x = self.act(self.fc11(x)) * self.fc12(x)\n x = self.drop(x)\n x = self.fc2(x)\n\n return x\n\n\ndef window_partition(x, window_size):\n \"\"\"\n Args:\n x: (B, H, W, C)\n window_size (int): window size\n\n Returns:\n windows: (num_windows*B, window_size, window_size, C)\n \"\"\"\n B, H, W, C = x.shape\n x = x.view(B, H // window_size, window_size, W // window_size, window_size, C)\n windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, C)\n return windows\n\n\ndef window_reverse(windows, window_size, H, W):\n \"\"\"\n Args:\n windows: (num_windows*B, window_size, window_size, C)\n window_size (int): Window size\n H (int): Height of image\n W (int): Width of image\n\n Returns:\n x: (B, H, W, C)\n \"\"\"\n B = int(windows.shape[0] / (H * W / window_size / window_size))\n x = windows.view(B, H // window_size, W // window_size, window_size, window_size, -1)\n x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(B, H, W, -1)\n return x\n\n\nclass WindowAttention(nn.Module):\n r\"\"\" Window based multi-head self attention (W-MSA) module with relative position bias.\n It supports both of shifted and non-shifted window.\n\n Args:\n dim (int): Number of input channels.\n window_size (tuple[int]): The height and width of the window.\n num_heads (int): Number of attention heads.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set\n attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0\n proj_drop (float, optional): Dropout ratio of output. Default: 0.0\n \"\"\"\n\n def __init__(self, dim, window_size, num_heads, qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):\n\n super().__init__()\n self.dim = dim\n self.window_size = window_size # Wh, Ww\n self.num_heads = num_heads\n head_dim = dim // num_heads\n self.scale = qk_scale or head_dim ** -0.5\n\n # define a parameter table of relative position bias\n self.relative_position_bias_table = nn.Parameter(\n torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)) # 2*Wh-1 * 2*Ww-1, nH\n\n # get pair-wise relative position index for each token inside the window\n coords_h = torch.arange(self.window_size[0])\n coords_w = torch.arange(self.window_size[1])\n coords = torch.stack(torch.meshgrid([coords_h, coords_w])) # 2, Wh, Ww\n coords_flatten = torch.flatten(coords, 1) # 2, Wh*Ww\n relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] # 2, Wh*Ww, Wh*Ww\n relative_coords = relative_coords.permute(1, 2, 0).contiguous() # Wh*Ww, Wh*Ww, 2\n relative_coords[:, :, 0] += self.window_size[0] - 1 # shift to start from 0\n relative_coords[:, :, 1] += self.window_size[1] - 1\n relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1\n relative_position_index = relative_coords.sum(-1) # Wh*Ww, Wh*Ww\n self.register_buffer(\"relative_position_index\", relative_position_index)\n\n self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n self.attn_drop = nn.Dropout(attn_drop)\n self.proj = nn.Linear(dim, dim)\n\n self.proj_drop = nn.Dropout(proj_drop)\n\n trunc_normal_(self.relative_position_bias_table, std=.02)\n self.softmax = nn.Softmax(dim=-1)\n\n def forward(self, x, mask=None):\n \"\"\"\n Args:\n x: input features with shape of (num_windows*B, N, C)\n mask: (0/-inf) mask with shape of (num_windows, Wh*Ww, Wh*Ww) or None\n \"\"\"\n\n B_, N, C = x.shape\n qkv = self.qkv(x).reshape(B_, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)\n q, k, v = qkv[0], qkv[1], qkv[2] # make torchscript happy (cannot use tensor as tuple)\n\n q = q * self.scale\n attn = (q @ k.transpose(-2, -1))\n\n relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(\n self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1) # Wh*Ww,Wh*Ww,nH\n relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() # nH, Wh*Ww, Wh*Ww\n attn = attn + relative_position_bias.unsqueeze(0)\n\n if mask is not None:\n nW = mask.shape[0]\n attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)\n attn = attn.view(-1, self.num_heads, N, N)\n attn = self.softmax(attn)\n else:\n attn = self.softmax(attn)\n\n attn = self.attn_drop(attn)\n\n x = (attn @ v).transpose(1, 2).reshape(B_, N, C)\n x = self.proj(x)\n x = self.proj_drop(x)\n\n return x\n\n def extra_repr(self) -> str:\n return f'dim={self.dim}, window_size={self.window_size}, num_heads={self.num_heads}'\n\n def flops(self, N):\n # calculate flops for 1 window with token length of N\n flops = 0\n # qkv = self.qkv(x)\n flops += N * self.dim * 3 * self.dim\n # attn = (q @ k.transpose(-2, -1))\n flops += self.num_heads * N * (self.dim // self.num_heads) * N\n # x = (attn @ v)\n flops += self.num_heads * N * N * (self.dim // self.num_heads)\n # x = self.proj(x)\n flops += N * self.dim * self.dim\n return flops\n\n@lru_cache()\ndef calculate_mask(x_size, window_size, shift_size):\n # calculate attention mask for SW-MSA\n H, W = x_size\n img_mask = torch.zeros((1, H, W, 1)) # 1 H W 1\n h_slices = (slice(0, -window_size),\n slice(-window_size, -shift_size),\n slice(-shift_size, None))\n w_slices = (slice(0, -window_size),\n slice(-window_size, -shift_size),\n slice(-shift_size, None))\n cnt = 0\n for h in h_slices:\n for w in w_slices:\n img_mask[:, h, w, :] = cnt\n cnt += 1\n\n mask_windows = window_partition(img_mask, window_size) # nW, window_size, window_size, 1\n mask_windows = mask_windows.view(-1, window_size * window_size)\n attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n\n return attn_mask\n\n\nclass SwinTransformerBlock(nn.Module):\n r\"\"\" Swin Transformer Block.\n\n Args:\n dim (int): Number of input channels.\n input_resolution (tuple[int]): Input resulotion.\n num_heads (int): Number of attention heads.\n window_size (int): Window size.\n shift_size (int): Shift size for SW-MSA.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n drop (float, optional): Dropout rate. Default: 0.0\n attn_drop (float, optional): Attention dropout rate. Default: 0.0\n drop_path (float, optional): Stochastic depth rate. Default: 0.0\n act_layer (nn.Module, optional): Activation layer. Default: nn.GELU\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n \"\"\"\n\n def __init__(self, dim, input_resolution, num_heads, window_size=7, shift_size=0,\n mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., drop_path=0.,\n act_layer=nn.GELU, norm_layer=nn.LayerNorm, use_checkpoint=False):\n super().__init__()\n self.dim = dim\n self.input_resolution = input_resolution\n self.num_heads = num_heads\n self.window_size = window_size\n self.shift_size = shift_size\n self.mlp_ratio = mlp_ratio\n self.use_checkpoint = use_checkpoint\n if min(self.input_resolution) <= self.window_size:\n # if window size is larger than input resolution, we don't partition windows\n self.shift_size = 0\n self.window_size = min(self.input_resolution)\n assert 0 <= self.shift_size < self.window_size, \"shift_size must in 0-window_size\"\n\n self.norm1 = norm_layer(dim)\n self.attn = WindowAttention(\n dim, window_size=to_2tuple(self.window_size), num_heads=num_heads,\n qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)\n\n self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()\n self.norm2 = norm_layer(dim)\n mlp_hidden_dim = int(dim * mlp_ratio)\n self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)\n\n # if self.shift_size > 0:\n # attn_mask = self.calculate_mask(self.input_resolution)\n # else:\n # attn_mask = None\n\n # self.register_buffer(\"attn_mask\", attn_mask)\n\n\n def forward(self, x, x_size):\n H, W = x_size\n B, L, C = x.shape\n # assert L == H * W, \"input feature has wrong size\"\n\n # if self.input_resolution != x_size:\n # self.input_resolution = x_size\n # if self.attn_mask is not None:\n # self.attn_mask = self.calculate_mask(x_size).to(x.device)\n\n shortcut = x\n x = self.norm1(x)\n x = x.view(B, H, W, C)\n\n # cyclic shift\n if self.shift_size > 0:\n shifted_x = torch.roll(x, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))\n else:\n shifted_x = x\n\n # partition windows\n x_windows = window_partition(shifted_x, self.window_size) # nW*B, window_size, window_size, C\n x_windows = x_windows.view(-1, self.window_size * self.window_size, C) # nW*B, window_size*window_size, C\n\n # W-MSA/SW-MSA (to be compatible for testing on images whose shapes are the multiple of window size\n\n # if self.input_resolution == x_size:\n # if self.use_checkpoint:\n # attn_windows = checkpoint.apply(self.attn, x_windows, self.attn_mask) # nW*B, window_size*window_size, C\n # else:\n # attn_windows = self.attn(x_windows, mask=self.attn_mask) # nW*B, window_size*window_size, C\n # else:\n # if self.use_checkpoint:\n # attn_windows = checkpoint.apply(self.attn, x_windows, self.calculate_mask(x_size).to(x.device))\n # else:\n # attn_windows = self.attn(x_windows, mask=self.calculate_mask(x_size).to(x.device))\n\n attn_mask = calculate_mask(x_size, self.window_size, self.shift_size).to(x.device)\n attn_windows = self.attn(x_windows, mask=attn_mask)\n\n # merge windows\n attn_windows = attn_windows.view(-1, self.window_size, self.window_size, C)\n shifted_x = window_reverse(attn_windows, self.window_size, H, W) # B H' W' C\n\n # reverse cyclic shift\n if self.shift_size > 0:\n x = torch.roll(shifted_x, shifts=(self.shift_size, self.shift_size), dims=(1, 2))\n else:\n x = shifted_x\n x = x.view(B, H * W, C)\n\n # FFN\n x = shortcut + self.drop_path(x)\n x = x + self.drop_path(self.mlp(self.norm2(x)))\n\n return x\n\n def extra_repr(self) -> str:\n return f\"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, \" \\\n f\"window_size={self.window_size}, shift_size={self.shift_size}, mlp_ratio={self.mlp_ratio}\"\n\n def flops(self):\n flops = 0\n H, W = self.input_resolution\n # norm1\n flops += self.dim * H * W\n # W-MSA/SW-MSA\n nW = H * W / self.window_size / self.window_size\n flops += nW * self.attn.flops(self.window_size * self.window_size)\n # mlp\n flops += 2 * H * W * self.dim * self.dim * self.mlp_ratio\n # norm2\n flops += self.dim * H * W\n return flops\n\n\nclass PatchMerging(nn.Module):\n r\"\"\" Patch Merging Layer.\n\n Args:\n input_resolution (tuple[int]): Resolution of input feature.\n dim (int): Number of input channels.\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n \"\"\"\n\n def __init__(self, input_resolution, dim, norm_layer=nn.LayerNorm):\n super().__init__()\n self.input_resolution = input_resolution\n self.dim = dim\n self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n self.norm = norm_layer(4 * dim)\n\n def forward(self, x):\n \"\"\"\n x: B, H*W, C\n \"\"\"\n H, W = self.input_resolution\n B, L, C = x.shape\n assert L == H * W, \"input feature has wrong size\"\n assert H % 2 == 0 and W % 2 == 0, f\"x size ({H}*{W}) are not even.\"\n\n x = x.view(B, H, W, C)\n\n x0 = x[:, 0::2, 0::2, :] # B H/2 W/2 C\n x1 = x[:, 1::2, 0::2, :] # B H/2 W/2 C\n x2 = x[:, 0::2, 1::2, :] # B H/2 W/2 C\n x3 = x[:, 1::2, 1::2, :] # B H/2 W/2 C\n x = torch.cat([x0, x1, x2, x3], -1) # B H/2 W/2 4*C\n x = x.view(B, -1, 4 * C) # B H/2*W/2 4*C\n\n x = self.norm(x)\n x = self.reduction(x)\n\n return x\n\n def extra_repr(self) -> str:\n return f\"input_resolution={self.input_resolution}, dim={self.dim}\"\n\n def flops(self):\n H, W = self.input_resolution\n flops = H * W * self.dim\n flops += (H // 2) * (W // 2) * 4 * self.dim * 2 * self.dim\n return flops\n\n\nclass BasicLayer(nn.Module):\n \"\"\" A basic Swin Transformer layer for one stage.\n\n Args:\n dim (int): Number of input channels.\n input_resolution (tuple[int]): Input resolution.\n depth (int): Number of blocks.\n num_heads (int): Number of attention heads.\n window_size (int): Local window size.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n drop (float, optional): Dropout rate. Default: 0.0\n attn_drop (float, optional): Attention dropout rate. Default: 0.0\n drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n \"\"\"\n\n def __init__(self, dim, input_resolution, depth, num_heads, window_size,\n mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,\n drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False):\n\n super().__init__()\n self.dim = dim\n self.input_resolution = input_resolution\n self.depth = depth\n self.use_checkpoint = False\n\n # build blocks\n self.blocks = nn.ModuleList([\n SwinTransformerBlock(dim=dim, input_resolution=input_resolution,\n num_heads=num_heads, window_size=window_size,\n shift_size=0 if (i % 2 == 0) else window_size // 2,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop, attn_drop=attn_drop,\n drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n norm_layer=norm_layer, use_checkpoint=use_checkpoint)\n for i in range(depth)])\n\n # patch merging layer\n if downsample is not None:\n self.downsample = downsample(input_resolution, dim=dim, norm_layer=norm_layer)\n else:\n self.downsample = None\n\n def forward(self, x, x_size):\n for i, blk in enumerate(self.blocks):\n if self.use_checkpoint:\n # x = checkpoint.checkpoint(blk, x, x_size)\n x = checkpoint.apply(blk, 2, x, x_size)\n else:\n x = blk(x, x_size)\n if self.downsample is not None:\n x = self.downsample(x)\n return x\n\n def extra_repr(self) -> str:\n return f\"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}\"\n\n def flops(self):\n flops = 0\n for blk in self.blocks:\n flops += blk.flops()\n if self.downsample is not None:\n flops += self.downsample.flops()\n return flops\n\n\nclass RSTB(nn.Module):\n \"\"\"Residual Swin Transformer Block (RSTB).\n\n Args:\n dim (int): Number of input channels.\n input_resolution (tuple[int]): Input resolution.\n depth (int): Number of blocks.\n num_heads (int): Number of attention heads.\n window_size (int): Local window size.\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim.\n qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set.\n drop (float, optional): Dropout rate. Default: 0.0\n attn_drop (float, optional): Attention dropout rate. Default: 0.0\n drop_path (float | tuple[float], optional): Stochastic depth rate. Default: 0.0\n norm_layer (nn.Module, optional): Normalization layer. Default: nn.LayerNorm\n downsample (nn.Module | None, optional): Downsample layer at the end of the layer. Default: None\n use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False.\n img_size: Input image size.\n patch_size: Patch size.\n resi_connection: The convolutional block before residual connection.\n \"\"\"\n\n def __init__(self, dim, input_resolution, depth, num_heads, window_size,\n mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0.,\n drop_path=0., norm_layer=nn.LayerNorm, downsample=None, use_checkpoint=False,\n img_size=224, patch_size=4, resi_connection='1conv'):\n super(RSTB, self).__init__()\n\n # print(f'dim: {dim}, input_resolution: {input_resolution}, depth: {depth}, num_heads: {num_heads}, window_size: {window_size}, img_size: {img_size}. patch_size: {patch_size}')\n\n self.dim = dim\n self.input_resolution = input_resolution\n\n self.residual_group = BasicLayer(dim=dim,\n input_resolution=input_resolution,\n depth=depth,\n num_heads=num_heads,\n window_size=window_size,\n mlp_ratio=mlp_ratio,\n qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop, attn_drop=attn_drop,\n drop_path=drop_path,\n norm_layer=norm_layer,\n downsample=downsample,\n use_checkpoint=use_checkpoint)\n\n if resi_connection == '1conv':\n self.conv = nn.Conv2d(dim, dim, 3, 1, 1)\n\n elif resi_connection == '3conv':\n # to save parameters and memory\n self.conv = nn.Sequential(nn.Conv2d(dim, dim // 4, 3, 1, 1), nn.LeakyReLU(negative_slope=0.2, inplace=True),\n nn.Conv2d(dim // 4, dim // 4, 1, 1, 0),\n nn.LeakyReLU(negative_slope=0.2, inplace=True),\n nn.Conv2d(dim // 4, dim, 3, 1, 1))\n\n self.patch_embed = PatchEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,\n norm_layer=None)\n\n self.patch_unembed = PatchUnEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=0, embed_dim=dim,\n norm_layer=None)\n\n def forward(self, x, x_size):\n x = self.patch_embed(self.conv(self.patch_unembed(self.residual_group(x, x_size), x_size))) + x\n return x\n\n\n def flops(self):\n flops = 0\n flops += self.residual_group.flops()\n H, W = self.input_resolution\n flops += H * W * self.dim * self.dim * 9\n flops += self.patch_embed.flops()\n flops += self.patch_unembed.flops()\n\n return flops\n\n\nclass PatchEmbed(nn.Module):\n r\"\"\" Image to Patch Embedding\n\n Args:\n img_size (int): Image size. Default: 224.\n patch_size (int): Patch token size. Default: 4.\n in_chans (int): Number of input image channels. Default: 3.\n embed_dim (int): Number of linear projection output channels. Default: 96.\n norm_layer (nn.Module, optional): Normalization layer. Default: None\n \"\"\"\n\n def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):\n super().__init__()\n img_size = to_2tuple(img_size)\n patch_size = to_2tuple(patch_size)\n patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]\n self.img_size = img_size\n self.patch_size = patch_size\n self.patches_resolution = patches_resolution\n self.num_patches = patches_resolution[0] * patches_resolution[1]\n\n self.in_chans = in_chans\n self.embed_dim = embed_dim\n\n if norm_layer is not None:\n self.norm = norm_layer(embed_dim)\n else:\n self.norm = None\n\n def forward(self, x, use_norm=True):\n x = x.flatten(2).transpose(1, 2) # B Ph*Pw C\n if use_norm and self.norm is not None:\n x = self.norm(x)\n return x\n\n def flops(self):\n flops = 0\n H, W = self.img_size\n if self.norm is not None:\n flops += H * W * self.embed_dim\n return flops\n\n\nclass PatchUnEmbed(nn.Module):\n r\"\"\" Image to Patch Unembedding\n\n Args:\n img_size (int): Image size. Default: 224.\n patch_size (int): Patch token size. Default: 4.\n in_chans (int): Number of input image channels. Default: 3.\n embed_dim (int): Number of linear projection output channels. Default: 96.\n norm_layer (nn.Module, optional): Normalization layer. Default: None\n \"\"\"\n\n def __init__(self, img_size=224, patch_size=4, in_chans=3, embed_dim=96, norm_layer=None):\n super().__init__()\n img_size = to_2tuple(img_size)\n patch_size = to_2tuple(patch_size)\n patches_resolution = [img_size[0] // patch_size[0], img_size[1] // patch_size[1]]\n self.img_size = img_size\n self.patch_size = patch_size\n self.patches_resolution = patches_resolution\n self.num_patches = patches_resolution[0] * patches_resolution[1]\n\n self.in_chans = in_chans\n self.embed_dim = embed_dim\n\n def forward(self, x, x_size):\n B, HW, C = x.shape\n x = x.transpose(1, 2).view(B, self.embed_dim, x_size[0], x_size[1]) # B Ph*Pw C\n return x\n\n def flops(self):\n flops = 0\n return flops\n\n\nclass Upsample(nn.Sequential):\n \"\"\"Upsample module.\n\n Args:\n scale (int): Scale factor. Supported scales: 2^n and 3.\n num_feat (int): Channel number of intermediate features.\n \"\"\"\n\n def __init__(self, scale, num_feat):\n m = []\n if (scale & (scale - 1)) == 0: # scale = 2^n\n for _ in range(int(math.log(scale, 2))):\n m.append(nn.Conv2d(num_feat, 4 * num_feat, 3, 1, 1))\n m.append(nn.PixelShuffle(2))\n elif scale == 3:\n m.append(nn.Conv2d(num_feat, 9 * num_feat, 3, 1, 1))\n m.append(nn.PixelShuffle(3))\n else:\n raise ValueError(f'scale {scale} is not supported. ' 'Supported scales: 2^n and 3.')\n super(Upsample, self).__init__(*m)\n\n\nclass UpsampleOneStep(nn.Sequential):\n \"\"\"UpsampleOneStep module (the difference with Upsample is that it always only has 1conv + 1pixelshuffle)\n Used in lightweight SR to save parameters.\n\n Args:\n scale (int): Scale factor. Supported scales: 2^n and 3.\n num_feat (int): Channel number of intermediate features.\n\n \"\"\"\n\n def __init__(self, scale, num_feat, num_out_ch, input_resolution=None):\n self.num_feat = num_feat\n self.input_resolution = input_resolution\n m = []\n m.append(nn.Conv2d(num_feat, (scale ** 2) * num_out_ch, 3, 1, 1))\n m.append(nn.PixelShuffle(scale))\n super(UpsampleOneStep, self).__init__(*m)\n\n def flops(self):\n H, W = self.input_resolution\n flops = H * W * self.num_feat * 3 * 9\n return flops\n\n\nclass SwinIR(nn.Module):\n r\"\"\" SwinIR\n A PyTorch impl of : `SwinIR: Image Restoration Using Swin Transformer`, based on Swin Transformer.\n\n Args:\n img_size (int | tuple(int)): Input image size. Default 64\n patch_size (int | tuple(int)): Patch size. Default: 1\n in_chans (int): Number of input image channels. Default: 3\n embed_dim (int): Patch embedding dimension. Default: 96\n depths (tuple(int)): Depth of each Swin Transformer layer.\n num_heads (tuple(int)): Number of attention heads in different layers.\n window_size (int): Window size. Default: 7\n mlp_ratio (float): Ratio of mlp hidden dim to embedding dim. Default: 4\n qkv_bias (bool): If True, add a learnable bias to query, key, value. Default: True\n qk_scale (float): Override default qk scale of head_dim ** -0.5 if set. Default: None\n drop_rate (float): Dropout rate. Default: 0\n attn_drop_rate (float): Attention dropout rate. Default: 0\n drop_path_rate (float): Stochastic depth rate. Default: 0.1\n norm_layer (nn.Module): Normalization layer. Default: nn.LayerNorm.\n ape (bool): If True, add absolute position embedding to the patch embedding. Default: False\n patch_norm (bool): If True, add normalization after patch embedding. Default: True\n use_checkpoint (bool): Whether to use checkpointing to save memory. Default: False\n upscale: Upscale factor. 2/3/4/8 for image SR, 1 for denoising and compress artifact reduction\n img_range: Image range. 1. or 255.\n upsampler: The reconstruction reconstruction module. 'pixelshuffle'/'pixelshuffledirect'/'nearest+conv'/None\n resi_connection: The convolutional block before residual connection. '1conv'/'3conv'\n \"\"\"\n\n def __init__(self, img_size=64, patch_size=1, in_chans=3,\n embed_dim=96, depths=[6, 6, 6, 6], num_heads=[6, 6, 6, 6],\n window_size=7, mlp_ratio=4., qkv_bias=True, qk_scale=None,\n drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,\n norm_layer=nn.LayerNorm, ape=False, patch_norm=True,\n use_checkpoint=False, upscale=2, img_range=1., upsampler='', resi_connection='1conv',\n **kwargs):\n super(SwinIR, self).__init__()\n num_in_ch = in_chans\n num_out_ch = in_chans\n num_feat = 64\n self.img_range = img_range\n if in_chans == 3:\n rgb_mean = (0.4488, 0.4371, 0.4040)\n self.mean = torch.Tensor(rgb_mean).view(1, 3, 1, 1)\n else:\n self.mean = torch.zeros(1, 1, 1, 1)\n self.upscale = upscale\n self.upsampler = upsampler\n self.window_size = window_size\n\n #####################################################################################################\n ################################### 1, shallow feature extraction ###################################\n self.conv_first = nn.Conv2d(num_in_ch, embed_dim, 3, 1, 1)\n\n #####################################################################################################\n ################################### 2, deep feature extraction ######################################\n self.num_layers = len(depths)\n self.embed_dim = embed_dim\n self.ape = ape\n self.patch_norm = patch_norm\n self.num_features = embed_dim\n self.mlp_ratio = mlp_ratio\n\n # split image into non-overlapping patches\n self.patch_embed = PatchEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,\n norm_layer=norm_layer if self.patch_norm else None)\n num_patches = self.patch_embed.num_patches\n patches_resolution = self.patch_embed.patches_resolution\n self.patches_resolution = patches_resolution\n # print('patches_resolution: ', patches_resolution)\n\n # merge non-overlapping patches into image\n self.patch_unembed = PatchUnEmbed(\n img_size=img_size, patch_size=patch_size, in_chans=embed_dim, embed_dim=embed_dim,\n norm_layer=norm_layer if self.patch_norm else None)\n\n # absolute position embedding\n if self.ape:\n self.absolute_pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))\n trunc_normal_(self.absolute_pos_embed, std=.02)\n\n self.pos_drop = nn.Dropout(p=drop_rate)\n\n # stochastic depth\n dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))] # stochastic depth decay rule\n\n # build Residual Swin Transformer blocks (RSTB)\n self.layers = nn.ModuleList()\n for i_layer in range(self.num_layers):\n layer = RSTB(dim=embed_dim,\n input_resolution=(patches_resolution[0],\n patches_resolution[1]),\n depth=depths[i_layer],\n num_heads=num_heads[i_layer],\n window_size=window_size,\n mlp_ratio=self.mlp_ratio,\n qkv_bias=qkv_bias, qk_scale=qk_scale,\n drop=drop_rate, attn_drop=attn_drop_rate,\n drop_path=dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])], # no impact on SR results\n norm_layer=norm_layer,\n downsample=None,\n use_checkpoint=use_checkpoint,\n img_size=img_size,\n patch_size=patch_size,\n resi_connection=resi_connection\n\n )\n self.layers.append(layer)\n self.norm = norm_layer(self.num_features)\n\n # build the last conv layer in deep feature extraction\n if resi_connection == '1conv':\n self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, 3, 1, 1)\n elif resi_connection == '3conv':\n # to save parameters and memory\n self.conv_after_body = nn.Sequential(nn.Conv2d(embed_dim, embed_dim // 4, 3, 1, 1),\n nn.LeakyReLU(negative_slope=0.2, inplace=True),\n nn.Conv2d(embed_dim // 4, embed_dim // 4, 1, 1, 0),\n nn.LeakyReLU(negative_slope=0.2, inplace=True),\n nn.Conv2d(embed_dim // 4, embed_dim, 3, 1, 1))\n\n #####################################################################################################\n ################################ 3, high quality image reconstruction ################################\n if self.upsampler == 'pixelshuffle':\n # for classical SR\n self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),\n nn.LeakyReLU(inplace=True))\n self.upsample = Upsample(upscale, num_feat)\n self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)\n elif self.upsampler == 'pixelshuffledirect':\n # for lightweight SR (to save parameters)\n self.upsample = UpsampleOneStep(upscale, embed_dim, num_out_ch,\n (patches_resolution[0], patches_resolution[1]))\n elif self.upsampler == 'nearest+conv':\n # for real-world SR (less artifacts)\n assert self.upscale == 4, 'only support x4 now.'\n self.conv_before_upsample = nn.Sequential(nn.Conv2d(embed_dim, num_feat, 3, 1, 1),\n nn.LeakyReLU(inplace=True))\n self.conv_up1 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)\n self.conv_up2 = nn.Conv2d(num_feat, num_feat, 3, 1, 1)\n self.conv_hr = nn.Conv2d(num_feat, num_feat, 3, 1, 1)\n self.conv_last = nn.Conv2d(num_feat, num_out_ch, 3, 1, 1)\n self.lrelu = nn.LeakyReLU(negative_slope=0.2, inplace=True)\n else:\n # for image denoising and JPEG compression artifact reduction\n self.conv_last = nn.Conv2d(embed_dim, num_out_ch, 3, 1, 1)\n\n self.apply(self._init_weights)\n\n def _init_weights(self, m):\n if isinstance(m, nn.Linear):\n trunc_normal_(m.weight, std=.02)\n if isinstance(m, nn.Linear) and m.bias is not None:\n nn.init.constant_(m.bias, 0)\n elif isinstance(m, nn.LayerNorm):\n nn.init.constant_(m.bias, 0)\n nn.init.constant_(m.weight, 1.0)\n\n @torch.jit.ignore\n def no_weight_decay(self):\n return {'absolute_pos_embed'}\n\n @torch.jit.ignore\n def no_weight_decay_keywords(self):\n return {'relative_position_bias_table'}\n\n def check_image_size(self, x):\n _, _, h, w = x.size()\n mod_pad_h = (self.window_size - h % self.window_size) % self.window_size\n mod_pad_w = (self.window_size - w % self.window_size) % self.window_size\n x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), 'reflect')\n return x\n\n def forward_features(self, x):\n x_size = (x.shape[2], x.shape[3])\n x = self.patch_embed(x)\n if self.ape:\n x = x + self.absolute_pos_embed\n x = self.pos_drop(x)\n\n for layer in self.layers:\n x = layer(x, x_size)\n\n x = self.norm(x) # B L C\n x = self.patch_unembed(x, x_size)\n\n return x\n\n def forward(self, x):\n H, W = x.shape[2:]\n x = self.check_image_size(x)\n \n self.mean = self.mean.type_as(x)\n x = (x - self.mean) * self.img_range\n\n if self.upsampler == 'pixelshuffle':\n # for classical SR\n x = self.conv_first(x)\n x = self.conv_after_body(self.forward_features(x)) + x\n x = self.conv_before_upsample(x)\n x = self.conv_last(self.upsample(x))\n elif self.upsampler == 'pixelshuffledirect':\n # for lightweight SR\n x = self.conv_first(x)\n x = self.conv_after_body(self.forward_features(x)) + x\n x = self.upsample(x)\n elif self.upsampler == 'nearest+conv':\n # for real-world SR\n x = self.conv_first(x)\n x = self.conv_after_body(self.forward_features(x)) + x\n x = self.conv_before_upsample(x)\n x = self.lrelu(self.conv_up1(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))\n x = self.lrelu(self.conv_up2(torch.nn.functional.interpolate(x, scale_factor=2, mode='nearest')))\n x = self.conv_last(self.lrelu(self.conv_hr(x)))\n else:\n # for image denoising and JPEG compression artifact reduction\n x_first = self.conv_first(x)\n res = self.conv_after_body(self.forward_features(x_first)) + x_first\n x = x + self.conv_last(res)\n\n x = x / self.img_range + self.mean\n\n return x[:, :, :H*self.upscale, :W*self.upscale]\n\n def flops(self):\n flops = 0\n H, W = self.patches_resolution\n flops += H * W * 3 * self.embed_dim * 9\n flops += self.patch_embed.flops()\n for i, layer in enumerate(self.layers):\n flops += layer.flops()\n flops += H * W * 3 * self.embed_dim * self.embed_dim\n flops += self.upsample.flops()\n return flops\n\n\nif __name__ == '__main__':\n upscale = 4\n window_size = 8\n height = (1024 // upscale // window_size + 1) * window_size\n width = (720 // upscale // window_size + 1) * window_size\n model = SwinIR(upscale=2, img_size=(height, width),\n window_size=window_size, img_range=1., depths=[6, 6, 6, 6],\n embed_dim=60, num_heads=[6, 6, 6, 6], mlp_ratio=2, upsampler='pixelshuffledirect')\n print(model)\n print(height, width, model.flops() / 1e9)\n\n x = torch.randn((1, 3, height, width))\n x = model(x)\n print(x.shape)\n"
},
{
"path": "code/real/bsrt/model/utils/interp_methods.py",
"content": "from math import pi\n\ntry:\n import torch\nexcept ImportError:\n torch = None\n\ntry:\n import numpy\nexcept ImportError:\n numpy = None\n\nif numpy is None and torch is None:\n raise ImportError(\"Must have either Numpy or PyTorch but both not found\")\n\n\ndef set_framework_dependencies(x):\n if type(x) is numpy.ndarray:\n to_dtype = lambda a: a\n fw = numpy\n else:\n to_dtype = lambda a: a.to(x.dtype)\n fw = torch\n eps = fw.finfo(fw.float32).eps\n return fw, to_dtype, eps\n\n\ndef support_sz(sz):\n def wrapper(f):\n f.support_sz = sz\n return f\n return wrapper\n\n@support_sz(4)\ndef cubic(x):\n fw, to_dtype, eps = set_framework_dependencies(x)\n absx = fw.abs(x)\n absx2 = absx ** 2\n absx3 = absx ** 3\n return ((1.5 * absx3 - 2.5 * absx2 + 1.) * to_dtype(absx <= 1.) +\n (-0.5 * absx3 + 2.5 * absx2 - 4. * absx + 2.) *\n to_dtype((1. < absx) & (absx <= 2.)))\n\n@support_sz(4)\ndef lanczos2(x):\n fw, to_dtype, eps = set_framework_dependencies(x)\n return (((fw.sin(pi * x) * fw.sin(pi * x / 2) + eps) /\n ((pi**2 * x**2 / 2) + eps)) * to_dtype(abs(x) < 2))\n\n@support_sz(6)\ndef lanczos3(x):\n fw, to_dtype, eps = set_framework_dependencies(x)\n return (((fw.sin(pi * x) * fw.sin(pi * x / 3) + eps) /\n ((pi**2 * x**2 / 3) + eps)) * to_dtype(abs(x) < 3))\n\n@support_sz(2)\ndef linear(x):\n fw, to_dtype, eps = set_framework_dependencies(x)\n return ((x + 1) * to_dtype((-1 <= x) & (x < 0)) + (1 - x) *\n to_dtype((0 <= x) & (x <= 1)))\n\n@support_sz(1)\ndef box(x):\n fw, to_dtype, eps = set_framework_dependencies(x)\n return to_dtype((-1 <= x) & (x < 0)) + to_dtype((0 <= x) & (x <= 1))\n"
},
{
"path": "code/real/bsrt/model/utils/psconv.py",
"content": "import torch\nimport torch.nn as nn\n\nclass PyConv2d(nn.Module):\n \"\"\"PyConv2d with padding (general case). Applies a 2D PyConv over an input signal composed of several input planes.\n Args:\n in_channels (int): Number of channels in the input image\n out_channels (list): Number of channels for each pyramid level produced by the convolution\n pyconv_kernels (list): Spatial size of the kernel for each pyramid level\n pyconv_groups (list): Number of blocked connections from input channels to output channels for each pyramid level\n stride (int or tuple, optional): Stride of the convolution. Default: 1\n dilation (int or tuple, optional): Spacing between kernel elements. Default: 1\n bias (bool, optional): If ``True``, adds a learnable bias to the output. Default: ``False``\n Example::\n >>> # PyConv with two pyramid levels, kernels: 3x3, 5x5\n >>> m = PyConv2d(in_channels=64, out_channels=[32, 32], pyconv_kernels=[3, 5], pyconv_groups=[1, 4])\n >>> input = torch.randn(4, 64, 56, 56)\n >>> output = m(input)\n >>> # PyConv with three pyramid levels, kernels: 3x3, 5x5, 7x7\n >>> m = PyConv2d(in_channels=64, out_channels=[16, 16, 32], pyconv_kernels=[3, 5, 7], pyconv_groups=[1, 4, 8])\n >>> input = torch.randn(4, 64, 56, 56)\n >>> output = m(input)\n \"\"\"\n def __init__(self, in_channels, out_channels, pyconv_kernels, pyconv_groups, stride=1, dilation=1, bias=False):\n super(PyConv2d, self).__init__()\n\n assert len(out_channels) == len(pyconv_kernels) == len(pyconv_groups)\n\n self.pyconv_levels = [None] * len(pyconv_kernels)\n for i in range(len(pyconv_kernels)):\n self.pyconv_levels[i] = nn.Conv2d(in_channels, out_channels[i], kernel_size=pyconv_kernels[i],\n stride=stride, padding=pyconv_kernels[i] // 2, groups=pyconv_groups[i],\n dilation=dilation, bias=bias)\n self.pyconv_levels = nn.ModuleList(self.pyconv_levels)\n\n def forward(self, x):\n out = []\n for level in self.pyconv_levels:\n out.append(level(x))\n\n return torch.cat(out, 1)\n\n################################################################\n\nclass PSConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, parts=4, bias=False):\n super(PSConv2d, self).__init__()\n self.gwconv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, dilation, dilation, groups=parts, bias=bias)\n self.gwconv_shift = nn.Conv2d(in_channels, out_channels, kernel_size, stride, 2 * dilation, 2 * dilation, groups=parts, bias=bias)\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, bias=bias)\n\n def backward_hook(grad):\n out = grad.clone()\n out[self.mask] = 0\n return out\n\n self.mask = torch.zeros(self.conv.weight.shape).byte().cuda()\n _in_channels = in_channels // parts\n _out_channels = out_channels // parts\n for i in range(parts):\n self.mask[i * _out_channels: (i + 1) * _out_channels, i * _in_channels: (i + 1) * _in_channels, : , :] = 1\n self.mask[(i + parts//2)%parts * _out_channels: ((i + parts//2)%parts + 1) * _out_channels, i * _in_channels: (i + 1) * _in_channels, :, :] = 1\n self.conv.weight.data[self.mask] = 0\n self.conv.weight.register_hook(backward_hook)\n\n self.weight = self.conv.weight\n self.bias = self.conv.bias\n\n def forward(self, x):\n x1, x2 = x.chunk(2, dim=1)\n x_shift = self.gwconv_shift(torch.cat((x2, x1), dim=1))\n return self.gwconv(x) + self.conv(x) + x_shift\n\n\n# PSConv-based Group Convolution\nclass PSGConv2d(nn.Module):\n def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1, dilation=1, groups=1, parts=4, bias=False):\n super(PSGConv2d, self).__init__()\n self.gwconv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups=groups * parts, bias=bias)\n self.gwconv_shift = nn.Conv2d(in_channels, out_channels, kernel_size, stride, 2 * padding, 2 * dilation, groups=groups * parts, bias=bias)\n self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, groups=groups, bias=bias)\n\n def backward_hook(grad):\n out = grad.clone()\n out[self.mask] = 0\n return out\n\n self.mask = torch.zeros(self.conv.weight.shape).bool().cuda()\n _in_channels = in_channels // (groups * parts)\n _out_channels = out_channels // (groups * parts)\n for i in range(parts):\n for j in range(groups):\n self.mask[(i + j * groups) * _out_channels: (i + j * groups + 1) * _out_channels, i * _in_channels: (i + 1) * _in_channels, : , :] = 1\n self.mask[((i + parts // 2) % parts + j * groups) * _out_channels: ((i + parts // 2) % parts + j * groups + 1) * _out_channels, i * _in_channels: (i + 1) * _in_channels, :, :] = 1\n self.conv.weight.data[self.mask] = 0\n self.conv.weight.register_hook(backward_hook)\n self.groups = groups\n\n self.weight = self.conv.weight\n self.bias = self.conv.bias\n\n def forward(self, x):\n x_split = (z.chunk(2, dim=1) for z in x.chunk(self.groups, dim=1))\n x_merge = torch.cat(tuple(torch.cat((x2, x1), dim=1) for (x1, x2) in x_split), dim=1)\n x_shift = self.gwconv_shift(x_merge)\n gx = self.gwconv(x)\n cx = self.conv(x)\n # print(x.shape, gx.shape, cx.shape, x_merge.shape, x_shift.shape)\n return gx + cx + x_shift\n\n"
},
{
"path": "code/real/bsrt/model/utils/resize_right.py",
"content": "import warnings\nfrom math import ceil\nimport model.utils.interp_methods as interp_methods\n\n\nclass NoneClass:\n pass\n\ntry:\n import torch\n from torch import nn\n nnModuleWrapped = nn.Module\nexcept ImportError:\n warnings.warn('No PyTorch found, will work only with Numpy')\n torch = None\n nnModuleWrapped = NoneClass\n\ntry:\n import numpy\nexcept ImportError:\n warnings.warn('No Numpy found, will work only with PyTorch')\n numpy = None\n\n\nif numpy is None and torch is None:\n raise ImportError(\"Must have either Numpy or PyTorch but both not found\")\n\n\ndef resize(input, scale_factors=None, out_shape=None,\n interp_method=interp_methods.cubic, support_sz=None,\n antialiasing=True):\n # get properties of the input tensor\n in_shape, n_dims = input.shape, input.ndim\n\n # fw stands for framework that can be either numpy or torch,\n # determined by the input type\n fw = numpy if type(input) is numpy.ndarray else torch\n eps = fw.finfo(fw.float32).eps\n\n # set missing scale factors or output shapem one according to another,\n # scream if both missing\n scale_factors, out_shape = set_scale_and_out_sz(in_shape, out_shape,\n scale_factors, fw)\n\n # sort indices of dimensions according to scale of each dimension.\n # since we are going dim by dim this is efficient\n sorted_filtered_dims_and_scales = [(dim, scale_factors[dim])\n for dim in sorted(range(n_dims),\n key=lambda ind: scale_factors[ind])\n if scale_factors[dim] != 1.]\n\n # unless support size is specified by the user, it is an attribute\n # of the interpolation method\n if support_sz is None:\n support_sz = interp_method.support_sz\n\n # when using pytorch, we need to know what is the input tensor device\n if fw is torch:\n device = input.device\n\n # output begins identical to input and changes with each iteration\n output = input\n\n # iterate over dims\n for dim, scale_factor in sorted_filtered_dims_and_scales:\n\n # get 1d set of weights and fields of view for each output location\n # along this dim\n field_of_view, weights = prepare_weights_and_field_of_view_1d(\n dim, scale_factor, in_shape[dim], out_shape[dim], interp_method,\n support_sz, antialiasing, fw, eps, device)\n\n # multiply the weights by the values in the field of view and\n # aggreagate\n output = apply_weights(output, field_of_view, weights, dim, n_dims,\n fw)\n return output\n\n\nclass ResizeLayer(nnModuleWrapped):\n def __init__(self, in_shape, scale_factors=None, out_shape=None,\n interp_method=interp_methods.cubic, support_sz=None,\n antialiasing=True):\n super(ResizeLayer, self).__init__()\n\n # fw stands for framework, that can be either numpy or torch. since\n # this is a torch layer, only one option in this case.\n fw = torch\n eps = fw.finfo(fw.float32).eps\n\n # set missing scale factors or output shapem one according to another,\n # scream if both missing\n scale_factors, out_shape = set_scale_and_out_sz(in_shape, out_shape,\n scale_factors, fw)\n\n # unless support size is specified by the user, it is an attribute\n # of the interpolation method\n if support_sz is None:\n support_sz = interp_method.support_sz\n\n self.n_dims = len(in_shape)\n\n # sort indices of dimensions according to scale of each dimension.\n # since we are going dim by dim this is efficient\n self.sorted_filtered_dims_and_scales = [(dim, scale_factors[dim])\n for dim in\n sorted(range(self.n_dims),\n key=lambda ind:\n scale_factors[ind])\n if scale_factors[dim] != 1.]\n\n # iterate over dims\n field_of_view_list = []\n weights_list = []\n for dim, scale_factor in self.sorted_filtered_dims_and_scales:\n\n # get 1d set of weights and fields of view for each output\n # location along this dim\n field_of_view, weights = prepare_weights_and_field_of_view_1d(\n dim, scale_factor, in_shape[dim], out_shape[dim],\n interp_method, support_sz, antialiasing, fw, eps, input.device)\n\n # keep weights and fields of views for all dims\n weights_list.append(nn.Parameter(weights, requires_grad=False))\n field_of_view_list.append(nn.Parameter(field_of_view,\n requires_grad=False))\n\n self.field_of_view = nn.ParameterList(field_of_view_list)\n self.weights = nn.ParameterList(weights_list)\n self.in_shape = in_shape\n\n def forward(self, input):\n # output begins identical to input and changes with each iteration\n output = input\n\n for (dim, scale_factor), field_of_view, weights in zip(\n self.sorted_filtered_dims_and_scales,\n self.field_of_view,\n self.weights):\n # multiply the weights by the values in the field of view and\n # aggreagate\n output = apply_weights(output, field_of_view, weights, dim,\n self.n_dims, torch)\n return output\n\n\ndef prepare_weights_and_field_of_view_1d(dim, scale_factor, in_sz, out_sz,\n interp_method, support_sz,\n antialiasing, fw, eps, device=None):\n # If antialiasing is taking place, we modify the window size and the\n # interpolation method (see inside function)\n interp_method, cur_support_sz = apply_antialiasing_if_needed(\n interp_method,\n support_sz,\n scale_factor,\n antialiasing)\n\n # STEP 1- PROJECTED GRID: The non-integer locations of the projection of\n # output pixel locations to the input tensor\n projected_grid = get_projected_grid(in_sz, out_sz, scale_factor, fw, device)\n\n # STEP 2- FIELDS OF VIEW: for each output pixels, map the input pixels\n # that influence it\n field_of_view = get_field_of_view(projected_grid, cur_support_sz, in_sz,\n fw, eps)\n\n # STEP 3- CALCULATE WEIGHTS: Match a set of weights to the pixels in the\n # field of view for each output pixel\n weights = get_weights(interp_method, projected_grid, field_of_view)\n\n return field_of_view, weights\n\n\ndef apply_weights(input, field_of_view, weights, dim, n_dims, fw):\n # STEP 4- APPLY WEIGHTS: Each output pixel is calculated by multiplying\n # its set of weights with the pixel values in its field of view.\n # We now multiply the fields of view with their matching weights.\n # We do this by tensor multiplication and broadcasting.\n # this step is separated to a different function, so that it can be\n # repeated with the same calculated weights and fields.\n\n # for this operations we assume the resized dim is the first one.\n # so we transpose and will transpose back after multiplying\n tmp_input = fw_swapaxes(input, dim, 0, fw)\n\n # field_of_view is a tensor of order 2: for each output (1d location\n # along cur dim)- a list of 1d neighbors locations.\n # note that this whole operations is applied to each dim separately,\n # this is why it is all in 1d.\n # neighbors = tmp_input[field_of_view] is a tensor of order image_dims+1:\n # for each output pixel (this time indicated in all dims), these are the\n # values of the neighbors in the 1d field of view. note that we only\n # consider neighbors along the current dim, but such set exists for every\n # multi-dim location, hence the final tensor order is image_dims+1.\n neighbors = tmp_input[field_of_view]\n\n # weights is an order 2 tensor: for each output location along 1d- a list\n # of weighs matching the field of view. we augment it with ones, for\n # broadcasting, so that when multiplies some tensor the weights affect\n # only its first dim.\n tmp_weights = fw.reshape(weights, (*weights.shape, * [1] * (n_dims - 1)))\n\n # now we simply multiply the weights with the neighbors, and then sum\n # along the field of view, to get a single value per out pixel\n tmp_output = (neighbors * tmp_weights).sum(1)\n\n # we transpose back the resized dim to its original position\n return fw_swapaxes(tmp_output, 0, dim, fw)\n\n\ndef set_scale_and_out_sz(in_shape, out_shape, scale_factors, fw):\n # eventually we must have both scale-factors and out-sizes for all in/out\n # dims. however, we support many possible partial arguments\n if scale_factors is None and out_shape is None:\n raise ValueError(\"either scale_factors or out_shape should be \"\n \"provided\")\n if out_shape is not None:\n # if out_shape has less dims than in_shape, we defaultly resize the\n # first dims for numpy and last dims for torch\n out_shape = (list(out_shape) + list(in_shape[:-len(out_shape)])\n if fw is numpy\n else list(in_shape[:-len(out_shape)]) + list(out_shape))\n if scale_factors is None:\n # if no scale given, we calculate it as the out to in ratio\n # (not recomended)\n scale_factors = [out_sz / in_sz for out_sz, in_sz\n in zip(out_shape, in_shape)]\n if scale_factors is not None:\n # by default, if a single number is given as scale, we assume resizing\n # two dims (most common are images with 2 spatial dims)\n scale_factors = (scale_factors\n if isinstance(scale_factors, (list, tuple))\n else [scale_factors, scale_factors])\n # if less scale_factors than in_shape dims, we defaultly resize the\n # first dims for numpy and last dims for torch\n scale_factors = (list(scale_factors) + [1] *\n (len(in_shape) - len(scale_factors)) if fw is numpy\n else [1] * (len(in_shape) - len(scale_factors)) +\n list(scale_factors))\n if out_shape is None:\n # when no out_shape given, it is calculated by multiplying the\n # scale by the in_shape (not recomended)\n out_shape = [ceil(scale_factor * in_sz)\n for scale_factor, in_sz in\n zip(scale_factors, in_shape)]\n # next line intentionally after out_shape determined for stability\n scale_factors = [float(sf) for sf in scale_factors]\n return scale_factors, out_shape\n\n\ndef get_projected_grid(in_sz, out_sz, scale_factor, fw, device=None):\n # we start by having the ouput coordinates which are just integer locations\n out_coordinates = fw.arange(out_sz)\n\n # if using torch we need to match the grid tensor device to the input device\n out_coordinates = fw_set_device(out_coordinates, device, fw)\n\n # This is projecting the ouput pixel locations in 1d to the input tensor,\n # as non-integer locations.\n # the following fomrula is derived in the paper\n # \"From Discrete to Continuous Convolutions\" by Shocher et al.\n return (out_coordinates / scale_factor +\n (in_sz - 1) / 2 - (out_sz - 1) / (2 * scale_factor))\n\n\ndef get_field_of_view(projected_grid, cur_support_sz, in_sz, fw, eps):\n # for each output pixel, map which input pixels influence it, in 1d.\n # we start by calculating the leftmost neighbor, using half of the window\n # size (eps is for when boundary is exact int)\n left_boundaries = fw_ceil(projected_grid - cur_support_sz / 2 - eps, fw)\n\n # then we simply take all the pixel centers in the field by counting\n # window size pixels from the left boundary\n ordinal_numbers = fw.arange(ceil(cur_support_sz - eps))\n # in case using torch we need to match the device\n ordinal_numbers = fw_set_device(ordinal_numbers, projected_grid.device, fw)\n field_of_view = left_boundaries[:, None] + ordinal_numbers\n\n # next we do a trick instead of padding, we map the field of view so that\n # it would be like mirror padding, without actually padding\n # (which would require enlarging the input tensor)\n mirror = fw_cat((fw.arange(in_sz), fw.arange(in_sz - 1, -1, step=-1)), fw)\n field_of_view = mirror[fw.remainder(field_of_view, mirror.shape[0])]\n field_of_view = fw_set_device(field_of_view,projected_grid.device, fw)\n return field_of_view\n\n\ndef get_weights(interp_method, projected_grid, field_of_view):\n # the set of weights per each output pixels is the result of the chosen\n # interpolation method applied to the distances between projected grid\n # locations and the pixel-centers in the field of view (distances are\n # directed, can be positive or negative)\n weights = interp_method(projected_grid[:, None] - field_of_view)\n\n # we now carefully normalize the weights to sum to 1 per each output pixel\n sum_weights = weights.sum(1, keepdims=True)\n sum_weights[sum_weights == 0] = 1\n return weights / sum_weights\n\n\ndef apply_antialiasing_if_needed(interp_method, support_sz, scale_factor,\n antialiasing):\n # antialiasing is \"stretching\" the field of view according to the scale\n # factor (only for downscaling). this is low-pass filtering. this\n # requires modifying both the interpolation (stretching the 1d\n # function and multiplying by the scale-factor) and the window size.\n if scale_factor >= 1.0 or not antialiasing:\n return interp_method, support_sz\n cur_interp_method = (lambda arg: scale_factor *\n interp_method(scale_factor * arg))\n cur_support_sz = support_sz / scale_factor\n return cur_interp_method, cur_support_sz\n\n\ndef fw_ceil(x, fw):\n if fw is numpy:\n return fw.int_(fw.ceil(x))\n else:\n return x.ceil().long()\n\n\ndef fw_cat(x, fw):\n if fw is numpy:\n return fw.concatenate(x)\n else:\n return fw.cat(x)\n\n\ndef fw_swapaxes(x, ax_1, ax_2, fw):\n if fw is numpy:\n return fw.swapaxes(x, ax_1, ax_2)\n else:\n return x.transpose(ax_1, ax_2)\n\ndef fw_set_device(x, device, fw):\n if fw is numpy:\n return x\n else:\n return x.to(device)\n"
},
{
"path": "code/real/bsrt/option.py",
"content": "import argparse\n\nparser = argparse.ArgumentParser(description='EDSR and MDSR')\n\nparser.add_argument('--n_resblocks', type=int, default=16,\n help='number of residual blocks')\nparser.add_argument('--n_feats', type=int, default=64,\n help='number of feature maps')\nparser.add_argument('--n_colors', type=int, default=3,\n help='number of color channels to use')\nparser.add_argument('--lr', type=float, default=1e-4,\n help='learning rate')\nparser.add_argument('--burst_size', type=int, default=14,\n help='burst size, max 14')\nparser.add_argument('--burst_channel', type=int, default=4,\n help='RAW channel, default:4')\nparser.add_argument('--swinfeature', action='store_true',\n help='use swin transformer to extract features')\nparser.add_argument('--model_level', type=str, default='S',\n help='S: small, L: large')\n\n################## fine-tune ##################\nparser.add_argument('--finetune', action='store_true',\n help='finetune model')\nparser.add_argument('--finetune_align', action='store_true',\n help='finetune alignment module')\nparser.add_argument('--finetune_swin', action='store_true',\n help='finetune swin trans module')\nparser.add_argument('--finetune_conv', action='store_true',\n help='finetune rest convs')\nparser.add_argument('--finetune_prelayer', action='store_true',\n help='finetune finetune pre feature extract layer')\nparser.add_argument('--finetune_upconv', action='store_true',\n help='finetune finetune up conv layer')\nparser.add_argument('--finetune_spynet', action='store_true',\n help='finetune finetune up conv layer')\n\n\n# Hardware specifications\nparser.add_argument('--n_threads', type=int, default=6,\n help='number of threads for data loading')\nparser.add_argument('--cpu', action='store_true',\n help='use cpu only')\nparser.add_argument('--n_GPUs', type=int, default=1,\n help='number of GPUs')\nparser.add_argument('--seed', type=int, default=1,\n help='random seed')\nparser.add_argument('--local_rank', type=int, default=-1,\n help='proc index')\nparser.add_argument('--fp16', action='store_true',\n help='use fp16 only')\nparser.add_argument('--use_checkpoint', action='store_true',\n help='use use_checkpoint in swin transformer')\n\n# Data specifications\nparser.add_argument('--root', type=str, default='/data/dataset/ntire21/burstsr/real',\n help='dataset directory')\nparser.add_argument('--val_root', type=str, default='../test_set',\n help='dataset directory')\nparser.add_argument('--mode', type=str, default='train',\n help='demo image directory')\nparser.add_argument('--scale', type=str, default='4',\n help='super resolution scale')\nparser.add_argument('--patch_size', type=int, default=256,\n help='output patch size')\nparser.add_argument('--rgb_range', type=int, default=1,\n help='maximum value of RGB')\n\nparser.add_argument('--chop', action='store_true',\n help='enable memory-efficient forward')\nparser.add_argument('--no_augment', action='store_true',\n help='do not use data augmentation')\n\n# Model specifications\nparser.add_argument('--model', default='LRSC_EDVR',\n help='model name')\n\nparser.add_argument('--act', type=str, default='relu',\n help='activation function')\nparser.add_argument('--pre_train', type=str, default='',\n help='pre-trained model directory')\nparser.add_argument('--extend', type=str, default='.',\n help='pre-trained model directory')\n\nparser.add_argument('--res_scale', type=float, default=1,\n help='residual scaling')\nparser.add_argument('--shift_mean', default=True,\n help='subtract pixel mean from the input')\nparser.add_argument('--dilation', action='store_true',\n help='use dilated convolution')\nparser.add_argument('--precision', type=str, default='single',\n choices=('single', 'half'),\n help='FP precision for test (single | half)')\n\n\n# Option for Residual channel attention network (RCAN)\nparser.add_argument('--n_resgroups', type=int, default=20,\n help='number of residual groups')\nparser.add_argument('--reduction', type=int, default=16,\n help='number of feature maps reduction')\nparser.add_argument('--DA', action='store_true',\n help='use Dual Attention')\nparser.add_argument('--CA', action='store_true',\n help='use Channel Attention')\nparser.add_argument('--non_local', action='store_true',\n help='use Dual Attention')\n\n# Training specifications\nparser.add_argument('--reset', action='store_true',\n help='reset the training')\nparser.add_argument('--test_every', type=int, default=1000,\n help='do test per every N batches')\nparser.add_argument('--epochs', type=int, default=100,\n help='number of epochs to train')\nparser.add_argument('--batch_size', type=int, default=8,\n help='input batch size for training')\nparser.add_argument('--split_batch', type=int, default=1,\n help='split the batch into smaller chunks')\nparser.add_argument('--self_ensemble', action='store_true',\n help='use self-ensemble method for test')\nparser.add_argument('--test_only', action='store_true',\n help='set this option to test the model')\nparser.add_argument('--gan_k', type=int, default=1,\n help='k value for adversarial loss')\n\n# Optimization specifications\n\nparser.add_argument('--decay', type=str, default='40-80',\n help='learning rate decay type')\nparser.add_argument('--gamma', type=float, default=0.5,\n help='learning rate decay factor for step decay')\nparser.add_argument('--optimizer', default='ADAM',\n choices=('SGD', 'ADAM', 'RMSprop'),\n help='optimizer to use (SGD | ADAM | RMSprop)')\nparser.add_argument('--momentum', type=float, default=0.9,\n help='SGD momentum')\nparser.add_argument('--betas', type=tuple, default=(0.9, 0.999),\n help='ADAM beta')\nparser.add_argument('--epsilon', type=float, default=1e-8,\n help='ADAM epsilon for numerical stability')\nparser.add_argument('--weight_decay', type=float, default=0,\n help='weight decay')\nparser.add_argument('--gclip', type=float, default=0,\n help='gradient clipping threshold (0 = no clipping)')\n\n# Loss specifications\nparser.add_argument('--loss', type=str, default='1*L1',\n help='loss function configuration')\nparser.add_argument('--skip_threshold', type=float, default='1e8',\n help='skipping batch that has large error')\n\n# Log specifications\nparser.add_argument('--save', type=str, default='test',\n help='file name to save')\nparser.add_argument('--load', type=str, default='',\n help='file name to load')\nparser.add_argument('--resume', type=int, default=0,\n help='resume from specific checkpoint')\nparser.add_argument('--save_models', action='store_true',\n help='save all intermediate models')\nparser.add_argument('--print_every', type=int, default=10,\n help='how many batches to wait before logging training status')\nparser.add_argument('--save_results', action='store_true',\n help='save output results')\nparser.add_argument('--save_gt', action='store_true',\n help='save low-resolution and high-resolution images together')\n\nargs = parser.parse_args()\n\nargs.scale = list(map(lambda x: int(x), args.scale.split('+')))\n\nif args.epochs == 0:\n args.epochs = 1e8\n\nfor arg in vars(args):\n if vars(args)[arg] == 'True':\n vars(args)[arg] = True\n elif vars(args)[arg] == 'False':\n vars(args)[arg] = False\n\n"
},
{
"path": "code/real/bsrt/pwcnet/LICENSE",
"content": "GNU GENERAL PUBLIC LICENSE\n Version 3, 29 June 2007\n\n Copyright (C) 2007 Free Software Foundation, Inc. | | |