Full Code of jxgu1016/Gabor_CNN_PyTorch for AI

Repository: jxgu1016/Gabor_CNN_PyTorch
Branch: master
Commit: 4549ec94c563
Files: 19
Total size: 31.2 KB

Directory structure:
gitextract_wjpk366e/

├── .gitattributes
├── .gitignore
├── LICENSE
├── README.md
├── demo/
│   ├── main.py
│   ├── net_factory.py
│   └── utils.py
├── gcn/
│   ├── __init__.py
│   ├── csrc/
│   │   ├── GOF.h
│   │   ├── cpu/
│   │   │   ├── GOF_cpu.cpp
│   │   │   └── vision.h
│   │   ├── cuda/
│   │   │   ├── GOF_cuda.cu
│   │   │   └── vision.h
│   │   └── vision.cpp
│   └── layers/
│       ├── GConv.py
│       ├── __init__.py
│       └── gradtest.py
├── install.sh
└── setup.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitattributes
================================================
# Auto detect text files and perform LF normalization
* text=auto


================================================
FILE: .gitignore
================================================
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
*.egg-info/
.installed.cfg
*.egg

# PyInstaller
#  Usually these files are written by a python script from a template
#  before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
.hypothesis/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
target/

# Jupyter Notebook
.ipynb_checkpoints

# pyenv
.python-version

# celery beat schedule file
celerybeat-schedule

# SageMath parsed files
*.sage.py

# Environments
.env
.venv
env/
venv/
ENV/

# Spyder project settings
.spyderproject
.spyproject

# Rope project settings
.ropeproject

# mkdocs documentation
/site

# mypy
.mypy_cache/

.DS_Store
runs*
*.json
CIFAR/*

================================================
FILE: LICENSE
================================================
MIT License

Copyright (c) 2018 jxgu1016

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

================================================
FILE: README.md
================================================
# WACV2018/TIP: Gabor Convolutional Networks

Official PyTorch implementation of Gabor CNN. 
But all the results in the paper are based on [Torch 7](https://github.com/bczhangbczhang/Gabor-Convolutional-Networks).
These two implementations are sharing the same infrastructure level code.

## Requirements
- PyTorch 1.1.0 (earlier versions are not supported)
- torchvision
  
## Install

```
git clone https://github.com/jxgu1016/Gabor_CNN_PyTorch
cd Gabor_CNN_PyTorch
sh install.sh
```

## Run MNIST demo

```
cd demo
python main.py --model gcn (--gpu 0)
```

## Please cite:
@article{GaborCNNs, title={Gabor Convolutional Networks}, author={Luan, Shangzhen and chen, chen and Zhang, Baochang* and Han, jungong and Liu, Jianzhuang}, year={2018}, IEEE Trans. Image processing. }


================================================
FILE: demo/main.py
================================================
from __future__ import division
import os
import time
import argparse
import torch
from torchvision import datasets, transforms
import torch.optim as optim
import torch.nn as nn
from torch.optim.lr_scheduler import StepLR, MultiStepLR
import torch.nn.functional as F
from utils import accuracy, AverageMeter, save_checkpoint, visualize_graph, get_parameters_size
from torch.utils.tensorboard import SummaryWriter
from net_factory import get_network_fn


parser = argparse.ArgumentParser(description='PyTorch GCN MNIST Training')
parser.add_argument('--epochs', default=50, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 4)')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=128, type=int,
                    metavar='N', help='mini-batch size (default: 64)')
parser.add_argument('--lr', '--learning-rate', default=0.01, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--print-freq', '-p', default=100, type=int,
                    metavar='N', help='print frequency (default: 10)')
parser.add_argument('--resume', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--pretrained', default='', type=str, metavar='PATH',
                    help='path to pretrained checkpoint (default: none)')
parser.add_argument('--gpu', default=-1, type=int,
                    metavar='N', help='GPU device ID (default: -1)')
parser.add_argument('--dataset_dir', default='../../MNIST', type=str, metavar='PATH',
                    help='path to dataset (default: ../MNIST)')
parser.add_argument('--comment', default='', type=str, metavar='INFO',
                    help='Extra description for tensorboard')
parser.add_argument('--model', default='', type=str, metavar='NETWORK',
                    help='Network to train')
args = parser.parse_args()

use_cuda = (args.gpu >= 0) and torch.cuda.is_available()
best_prec1 = 0
writer = SummaryWriter(comment='_'+args.model+'_'+args.comment)
iteration = 0

# Prepare the MNIST dataset
normalize = transforms.Normalize((0.1307,), (0.3081,))
train_transform = transforms.Compose([
    transforms.ToTensor(),
    normalize,
    ])
test_transform = transforms.Compose([
    transforms.ToTensor(), 
    normalize,
    ])


train_dataset = datasets.MNIST(root=args.dataset_dir, train=True, 
    				download=True, transform=train_transform)
test_dataset = datasets.MNIST(root=args.dataset_dir, train=False, 
                    download=True,transform=test_transform)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=args.batch_size,
                                num_workers=args.workers, pin_memory=True, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=args.batch_size,
                                num_workers=args.workers, pin_memory=True, shuffle=True)

# Load model
model = get_network_fn(args.model)
print(model)

# Try to visulize the model
try:
	visualize_graph(model, writer, input_size=(1, 1, 28, 28))
except:
	print('\nNetwork Visualization Failed! But the training procedure continue.')

# optimizer = optim.Adadelta(model.parameters(), lr=args.lr, rho=0.9, eps=1e-06, weight_decay=3e-05)
# optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=3e-05)
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum, weight_decay=3e-05)
scheduler = StepLR(optimizer, step_size=10, gamma=0.5)
criterion = nn.CrossEntropyLoss()

device = torch.device("cuda" if use_cuda else "cpu")
model = model.to(device)
criterion = criterion.to(device)

# Calculate the total parameters of the model
print('Model size: {:0.2f} million float parameters'.format(get_parameters_size(model)/1e6))

if args.pretrained:
    if os.path.isfile(args.pretrained):
        print("=> loading checkpoint '{}'".format(args.pretrained))
        checkpoint = torch.load(args.pretrained)
        model.load_state_dict(checkpoint['state_dict'])
    else:
        print("=> no checkpoint found at '{}'".format(args.pretrained))

def train(epoch):
    model.train()
    global iteration
    st = time.time()
    for batch_idx, (data, target) in enumerate(train_loader):
        iteration += 1
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        prec1, = accuracy(output, target)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % args.print_freq == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}, Accuracy: {:.2f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item(), prec1.item()))
            writer.add_scalar('Loss/Train', loss.item(), iteration)
            writer.add_scalar('Accuracy/Train', prec1, iteration)
    epoch_time = time.time() - st
    print('Epoch time:{:0.2f}s'.format(epoch_time))
    scheduler.step()

def test(epoch):
    model.eval()
    test_loss = AverageMeter()
    acc = AverageMeter()
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss.update(F.cross_entropy(output, target, reduction='mean').item(), target.size(0))
            prec1, = accuracy(output, target) # test precison in one batch
            acc.update(prec1.item(), target.size(0))
    print('\nTest set: Average loss: {:.4f}, Accuracy: {:.2f}%\n'.format(test_loss.avg, acc.avg))
    writer.add_scalar('Loss/Test', test_loss.avg, epoch)
    writer.add_scalar('Accuracy/Test', acc.avg, epoch)
    return acc.avg

for epoch in range(args.start_epoch, args.epochs):
    print('------------------------------------------------------------------------')
    train(epoch+1)
    prec1 = test(epoch+1)

    # remember best prec@1 and save checkpoint
    is_best = prec1 > best_prec1
    best_prec1 = max(prec1, best_prec1)
    save_checkpoint({
        'epoch': epoch + 1,
        'state_dict': model.state_dict(),
        'best_prec1': best_prec1,
        'optimizer' : optimizer.state_dict(),
    }, is_best)

print('Finished!')
print('Best Test Precision@top1:{:.2f}'.format(best_prec1))
writer.add_scalar('Best TOP1', best_prec1, 0)
writer.close()

================================================
FILE: demo/net_factory.py
================================================
from __future__ import division
import torch
import torch.nn as nn
import torch.nn.functional as F
from gcn.layers import GConv

class GCN(nn.Module):
    def __init__(self, channel=4):
        super(GCN, self).__init__()
        self.channel = channel
        self.model = nn.Sequential(
            GConv(1, 10, 5, padding=2, stride=1, M=channel, nScale=1, bias=False, expand=True),
            nn.BatchNorm2d(10*channel),
            nn.ReLU(inplace=True),

            GConv(10, 20, 5, padding=2, stride=1, M=channel, nScale=2, bias=False),
            nn.BatchNorm2d(20*channel),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2,2),

            GConv(20, 40, 5, padding=0, stride=1, M=channel, nScale=3, bias=False),
            nn.BatchNorm2d(40*channel),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2,2),

            GConv(40, 80, 5, padding=0, stride=1, M=channel, nScale=4, bias=False),
            nn.BatchNorm2d(80*channel),
            nn.ReLU(inplace=True),
        )
        self.fc1 = nn.Linear(80, 1024)
        self.relu = nn.ReLU(inplace=True)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = self.model(x)
        # x = x.view(-1, self.channel, 80)
        # x = torch.max(x, 1)[0]
        ## x = x.view(-1, 80 * self.channel)
        x = x.view(-1, 80, self.channel)
        x = torch.max(x, 2)[0]
        x = self.fc1(x)
        x = self.relu(x)
        x = self.dropout(x)
        x = self.fc2(x)
        return x


def get_network_fn(name):
    networks_zoo = {
    'gcn': GCN(channel=4),
    }
    if name is '':
        raise ValueError('Specify the network to train. All networks available:{}'.format(networks_zoo.keys()))
    elif name not in networks_zoo:
        raise ValueError('Name of network unknown {}. All networks available:{}'.format(name, networks_zoo.keys()))
    return networks_zoo[name]

================================================
FILE: demo/utils.py
================================================
import shutil
import torch
from torch.utils.tensorboard import SummaryWriter


class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    with torch.no_grad():
        maxk = max(topk)
        batch_size = target.size(0)

        _, pred = output.topk(maxk, 1, True, True)
        pred = pred.t()
        correct = pred.eq(target.view(1, -1).expand_as(pred))

        res = []
        for k in topk:
            correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
            res.append(correct_k.mul_(100.0 / batch_size))
        return res

def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

def visualize_graph(model, writer, input_size=(1, 3, 32, 32)):
    dummy_input = torch.rand(input_size)
    # with SummaryWriter(comment=name) as w:
    writer.add_graph(model, (dummy_input, ))

def get_parameters_size(model):
    total = 0
    for p in model.parameters():
        _size = 1
        for i in range(len(p.size())):
            _size *= p.size(i)
        total += _size
    return total

================================================
FILE: gcn/__init__.py
================================================


================================================
FILE: gcn/csrc/GOF.h
================================================
#pragma once

#include "cpu/vision.h"

#ifdef WITH_CUDA
#include "cuda/vision.h"
#endif

// Interface for Python
at::Tensor GOF_forward(const at::Tensor& weight, 
                       const at::Tensor& gaborFilterBank) {
  if (weight.type().is_cuda()) {
#ifdef WITH_CUDA
    return GOF_forward_cuda(weight, gaborFilterBank);
#else
    AT_ERROR("Not compiled with GPU support");
#endif
  }
  return GOF_forward_cpu(weight, gaborFilterBank);
}

at::Tensor GOF_backward(const at::Tensor& grad_output,
                        const at::Tensor& gaborFilterBank) {
  if (grad_output.type().is_cuda()) {
#ifdef WITH_CUDA
    return GOF_backward_cuda(grad_output, gaborFilterBank);
#else
    AT_ERROR("Not compiled with GPU support");
#endif
  }
  return GOF_backward_cpu(grad_output, gaborFilterBank);
}



================================================
FILE: gcn/csrc/cpu/GOF_cpu.cpp
================================================
#include "cpu/vision.h"


template <typename T>
void GOFForward_cpu_kernel(
  const T* weight_data,
  const T* gaborFilterBank_data,
  const int nOutputPlane,
  const int nInputPlane,
  const int nChannel,
  const int kH,
  const int kW,
  T* output_data) {
  for (int i = 0; i < nOutputPlane; i++) {
    for (int j = 0; j < nInputPlane; j++) {
      for (int l = 0; l < nChannel * kH * kW; l++) {
        T val = *(weight_data + i * (nInputPlane * nChannel * kH * kW)
                              + j * (nChannel * kH * kW)
                              + l);
        for (int k = 0; k < nChannel; k++) {
          T gabortmp = *(gaborFilterBank_data + k * (kW * kH) 
                                              + l % (kW * kH));
          T *target = output_data + i * (nChannel * nInputPlane * nChannel * kH * kW)
                                  + k * (nInputPlane * nChannel * kH * kW)
                                  + j * (nChannel * kH * kW)
                                  + l;
          *target = val * gabortmp;
        }
      }
    }
  }
}

template <typename T>
void GOFBackward_cpu_kernel(
  const T* grad_output_data,
  const T* gaborFilterBank_data,
  const int nOutputPlane,
  const int nInputPlane,
  const int nChannel,
  const int kH,
  const int kW,
  T* grad_weight_data) {
  const int nEntry = nChannel * kH * kW;

  for (int i = 0; i < nOutputPlane; i++) {
    for (int j = 0; j < nInputPlane; j++) {
      for (int l = 0; l < nEntry; l++) {
        T *val = grad_weight_data + i * (nInputPlane * nEntry)
                                  + j * (nEntry) + l;
        *val = 0;
        for (int k = 0; k < nChannel; k++) {
          T gabortmp = *(gaborFilterBank_data + k * (kW * kH)
                                              + l % (kW * kH));
          T target = *(grad_output_data + i * (nChannel * nInputPlane * nEntry)
                                       + k * (nInputPlane * nEntry)
                                       + j * (nEntry)
                                       + l);
          *val = *val + target * gabortmp;
        }
      }
    }
  }
}


at::Tensor GOF_forward_cpu(const at::Tensor& weight,
                           const at::Tensor& gaborFilterBank) {
  AT_ASSERTM(!weight.type().is_cuda(), "weight must be a CPU tensor");
  AT_ASSERTM(!gaborFilterBank.type().is_cuda(), "gaborFilterBank must be a CPU tensor");

  auto nOutputPlane = weight.size(0);
  auto nInputPlane = weight.size(1);
  auto nChannel = weight.size(2);
  auto kH = weight.size(3);
  auto kW = weight.size(4);

  auto output = at::empty({nOutputPlane * nChannel, nInputPlane * nChannel, kH, kW}, weight.options());

  if (output.numel() == 0) {
    return output;
  }

  AT_DISPATCH_FLOATING_TYPES(weight.type(), "GOF_forward", [&] {
    GOFForward_cpu_kernel<scalar_t>(
         weight.data<scalar_t>(),
         gaborFilterBank.data<scalar_t>(),
         nOutputPlane,
         nInputPlane,
         nChannel,
         kH,
         kW,
         output.data<scalar_t>());
  });
  return output;
}

at::Tensor GOF_backward_cpu(const at::Tensor& grad_output,
                            const at::Tensor& gaborFilterBank) {
  AT_ASSERTM(!grad_output.type().is_cuda(), "grad_output must be a CPU tensor");
  AT_ASSERTM(!gaborFilterBank.type().is_cuda(), "gaborFilterBank must be a CPU tensor");

  auto nChannel = gaborFilterBank.size(0);
  auto nOutputPlane = grad_output.size(0) / nChannel;
  auto nInputPlane = grad_output.size(1) / nChannel;
  auto kH = grad_output.size(2);
  auto kW = grad_output.size(3);

  auto grad_weight = at::empty({nOutputPlane, nInputPlane, nChannel, kH, kW}, grad_output.options());

  if (grad_weight.numel() == 0) {
    return grad_weight;
  }

  AT_DISPATCH_FLOATING_TYPES(grad_output.type(), "GOF_backward", [&] {
    GOFBackward_cpu_kernel<scalar_t>(
         grad_output.data<scalar_t>(),
         gaborFilterBank.data<scalar_t>(),
         nOutputPlane,
         nInputPlane,
         nChannel,
         kH,
         kW,
         grad_weight.data<scalar_t>());
  });
  return grad_weight;
}

================================================
FILE: gcn/csrc/cpu/vision.h
================================================
#pragma once
#include <torch/extension.h>


at::Tensor GOF_forward_cpu(const at::Tensor& weight, 
                           const at::Tensor& gaborFilterBank);

at::Tensor GOF_backward_cpu(const at::Tensor& grad_output,
                            const at::Tensor& gaborFilterBank);


================================================
FILE: gcn/csrc/cuda/GOF_cuda.cu
================================================
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>

#include <THC/THC.h>
#include <THC/THCAtomics.cuh>
#include <THC/THCDeviceUtils.cuh>

// TODO make it in a common file
#define CUDA_1D_KERNEL_LOOP(i, n)                            \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; \
       i += blockDim.x * gridDim.x)


template <typename T>
__global__ void GOFForward_cuda_kernel(const int nthreads,
                                       const T* weight_data,
                                       const T* gaborFilterBank_data, 
                                       const int nOutputPlane,
                                       const int nInputPlane,
                                       const int nChannel,
                                       const int kH,
                                       const int kW,
                                       T* output_data) {
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
    auto w = index % kW;
    auto h = (index / kW) % kH;
    auto c = (index / kW / kH) % nChannel;
    auto in = (index / kW / kH / nChannel) % nInputPlane;
    auto ori = (index / kW / kH / nChannel / nInputPlane) % nChannel;
    auto ou = index / kW / kH / nChannel / nInputPlane / nChannel;
    T val = *(weight_data + (((ou * nInputPlane + in) * nChannel + c) * kH + h) * kW + w);
    T *target = output_data + index;
    T gabortmp = *(gaborFilterBank_data + ori * (kH * kW)
                                        + h * kW
                                        + w);
    *target = val * gabortmp;
  }
}

template <typename T>
__global__ void GOFBackward_cuda_kernel(const int nthreads,
                                       const T* grad_output_data,
                                       const T* gaborFilterBank_data, 
                                       const int nOutputPlane,
                                       const int nInputPlane,
                                       const int nChannel,
                                       const int kH,
                                       const int kW,
                                       T* grad_weight_data) {
  auto nEntry = nChannel * kH * kW;
  CUDA_1D_KERNEL_LOOP(index, nthreads) {
    auto l = index % nEntry;
    auto j = (index / nEntry) % nInputPlane;
    auto i = index / nEntry / nInputPlane;
    T *val = grad_weight_data + index;
    *val = 0;
    for (int k = 0; k < nChannel; k++) {
      T gabortmp = *(gaborFilterBank_data + k * (kW * kH)
                                          + l % (kW * kH));
      T target = *(grad_output_data + i * (nChannel * nInputPlane * nEntry)
                                    + k * (nInputPlane * nEntry)
                                    + j * (nEntry)
                                    + l);     
			*val = *val + target * gabortmp;
    }
  }
}

at::Tensor GOF_forward_cuda(const at::Tensor& weight,
                            const at::Tensor& gaborFilterBank) {
  AT_ASSERTM(weight.type().is_cuda(), "weight must be a CUDA tensor");
  AT_ASSERTM(gaborFilterBank.type().is_cuda(), "gaborFilterBank must be a CUDA tensor");

  auto nOutputPlane = weight.size(0);
  auto nInputPlane = weight.size(1);
  auto nChannel = weight.size(2);
  auto kH = weight.size(3);
  auto kW = weight.size(4);

  auto output = at::empty({nOutputPlane * nChannel, nInputPlane * nChannel, kH, kW}, weight.options());
  // auto nEntry = nChannel * kH * kW;
  auto output_size = nOutputPlane * nChannel* nInputPlane * nChannel * kH * kW;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  dim3 grid(std::min(THCCeilDiv(output_size, 512L), 4096L));
  dim3 block(512);

  if (output.numel() == 0) {
    THCudaCheck(cudaGetLastError());
    return output;
  }

  AT_DISPATCH_FLOATING_TYPES(weight.type(), "GOF_forward", [&] {
    GOFForward_cuda_kernel<scalar_t><<<grid, block, 0, stream>>>(
      output_size,
      weight.data<scalar_t>(),
      gaborFilterBank.data<scalar_t>(),
      nOutputPlane,
      nInputPlane,
      nChannel,
      kH,
      kW,
      output.data<scalar_t>());
  });
  THCudaCheck(cudaGetLastError());
  return output;
}

at::Tensor GOF_backward_cuda(const at::Tensor& grad_output,
                             const at::Tensor& gaborFilterBank) {
  AT_ASSERTM(grad_output.type().is_cuda(), "grad_output must be a CUDA tensor");
  AT_ASSERTM(gaborFilterBank.type().is_cuda(), "gaborFilterBank must be a CUDA tensor");

  auto nChannel = gaborFilterBank.size(0);
  auto nOutputPlane = grad_output.size(0) / nChannel;
  auto nInputPlane = grad_output.size(1) / nChannel;
  auto kH = grad_output.size(2);
  auto kW = grad_output.size(3);

  auto grad_weight = at::empty({nOutputPlane, nInputPlane, nChannel, kH, kW}, grad_output.options());
  auto nEntry = nChannel * kH * kW;
  auto grad_weight_size = nOutputPlane * nInputPlane * nEntry;
  cudaStream_t stream = at::cuda::getCurrentCUDAStream();

  dim3 grid(std::min(THCCeilDiv(grad_weight_size, 512L), 4096L));
  dim3 block(512);

  if (grad_weight.numel() == 0) {
    THCudaCheck(cudaGetLastError());
    return grad_weight;
  }

  AT_DISPATCH_FLOATING_TYPES(grad_output.type(), "GOF_backward", [&] {
    GOFBackward_cuda_kernel<scalar_t><<<grid, block, 0, stream>>>(
      grad_weight_size,
      grad_output.data<scalar_t>(),
      gaborFilterBank.data<scalar_t>(),
      nOutputPlane,
      nInputPlane,
      nChannel,
      kH,
      kW,
      grad_weight.data<scalar_t>());
  });
  THCudaCheck(cudaGetLastError());
  return grad_weight;
}

================================================
FILE: gcn/csrc/cuda/vision.h
================================================
#pragma once
#include <torch/extension.h>


at::Tensor GOF_forward_cuda(const at::Tensor& weight, 
                            const at::Tensor& gaborFilterBank);

at::Tensor GOF_backward_cuda(const at::Tensor& grad_output,
                             const at::Tensor& gaborFilterBank);

================================================
FILE: gcn/csrc/vision.cpp
================================================
#include "GOF.h"

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("gof_forward", &GOF_forward, "GOF forward");
  m.def("gof_backward", &GOF_backward, "GOF backward");
}


================================================
FILE: gcn/layers/GConv.py
================================================
from __future__ import division
import math
import torch
from torch import nn
import torch.nn.functional as F
from torch.autograd import Function
from torch.nn.modules.utils import _pair
from torch.nn.modules.conv import _ConvNd
from torch.autograd.function import once_differentiable

from gcn import _C

class GOF_Function(Function):
    @staticmethod
    def forward(ctx, weight, gaborFilterBank):
        ctx.save_for_backward(weight, gaborFilterBank)
        output = _C.gof_forward(weight, gaborFilterBank)
        return output

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        weight, gaborFilterBank = ctx.saved_tensors
        grad_weight = _C.gof_backward(grad_output, gaborFilterBank)
        return grad_weight, None 

class MConv(_ConvNd):
    '''
    Baee layer class for modulated convolution
    '''
    def __init__(self, in_channels, out_channels, kernel_size, M=4, nScale=3, stride=1,
                    padding=0, dilation=1, groups=1, bias=True, expand=False, padding_mode='zeros'):
        if groups != 1:
            raise ValueError('Group-conv not supported!')
        kernel_size = (M,) + _pair(kernel_size)
        stride = _pair(stride)
        padding = _pair(padding)
        dilation = _pair(dilation)
        super(MConv, self).__init__(
            in_channels, out_channels, kernel_size, stride, padding, dilation,
            False, _pair(0), groups, bias, padding_mode)
        self.expand = expand
        self.M = M
        self.need_bias = bias
        self.generate_MFilters(nScale, kernel_size)
        self.GOF_Function = GOF_Function.apply

    def generate_MFilters(self, nScale, kernel_size):
        raise NotImplementedError

    def forward(self, x):
        if self.expand:
            x = self.do_expanding(x)
        new_weight = self.GOF_Function(self.weight, self.MFilters)
        new_bias = self.expand_bias(self.bias) if self.need_bias else self.bias
        if self.padding_mode == 'circular':
            expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
            return F.conv2d(F.pad(input, expanded_padding, mode='circular'),
                            self.weight, self.bias, self.stride,
                            _pair(0), self.dilation, self.groups)
        return F.conv2d(x, new_weight, new_bias, self.stride,
                self.padding, self.dilation, self.groups)

    def do_expanding(self, x):
        index = []
        for i in range(x.size(1)):
            for _ in range(self.M):
                index.append(i)
        index = torch.LongTensor(index).cuda() if x.is_cuda else torch.LongTensor(index)
        return x.index_select(1, index)
    
    def expand_bias(self, bias):
        index = []
        for i in range(bias.size()):
            for _ in range(self.M):
                index.append(i)
        index = torch.LongTensor(index).cuda() if bias.is_cuda else torch.LongTensor(index)
        return bias.index_select(0, index)

class GConv(MConv):
    '''
    Gabor Convolutional Operation Layer
    '''
    def __init__(self, in_channels, out_channels, kernel_size, M=4, nScale=3, stride=1,
                    padding=0, dilation=1, groups=1, bias=True, expand=False, padding_mode='zeros'):
        super(GConv, self).__init__(in_channels, out_channels, kernel_size, M, nScale, stride,
                    padding, dilation, groups, bias, expand, padding_mode)

    def generate_MFilters(self, nScale, kernel_size):
        # To generate Gabor Filters
        self.register_buffer('MFilters', getGaborFilterBank(nScale, *kernel_size))

def getGaborFilterBank(nScale, M, h, w):
    Kmax = math.pi / 2
    f = math.sqrt(2)
    sigma = math.pi
    sqsigma = sigma ** 2
    postmean = math.exp(-sqsigma / 2)
    if h != 1:
        gfilter_real = torch.zeros(M, h, w)
        for i in range(M):
            theta = i / M * math.pi
            k = Kmax / f ** (nScale - 1)
            xymax = -1e309
            xymin = 1e309
            for y in range(h):
                for x in range(w):
                    y1 = y + 1 - ((h + 1) / 2)
                    x1 = x + 1 - ((w + 1) / 2)
                    tmp1 = math.exp(-(k * k * (x1 * x1 + y1 * y1) / (2 * sqsigma)))
                    tmp2 = math.cos(k * math.cos(theta) * x1 + k * math.sin(theta) * y1) - postmean # For real part
                    # tmp3 = math.sin(k*math.cos(theta)*x1+k*math.sin(theta)*y1) # For imaginary part
                    gfilter_real[i][y][x] = k * k * tmp1 * tmp2 / sqsigma			
                    xymax = max(xymax, gfilter_real[i][y][x])
                    xymin = min(xymin, gfilter_real[i][y][x])
            gfilter_real[i] = (gfilter_real[i] - xymin) / (xymax - xymin)
    else:
        gfilter_real = torch.ones(M, h, w)
    return gfilter_real

================================================
FILE: gcn/layers/__init__.py
================================================
from .GConv import GConv

__all__=['GConv']

================================================
FILE: gcn/layers/gradtest.py
================================================
import torch
from torch.autograd import gradcheck
from gcn.layers.GConv import GOF_Function

def gradchecking(use_cuda=False):
    print('-'*80)
    GOF = GOF_Function.apply
    device = torch.device("cuda" if use_cuda else "cpu")

    weight = torch.randn(8,8,4,3,3).to(device).double().requires_grad_()
    gfb = torch.randn(4,3,3).to(device).double()

    test = gradcheck(GOF, (weight, gfb), eps=1e-6, atol=1e-4, rtol=1e-3, raise_exception=True)
    print(test)


if __name__ == "__main__":
    gradchecking()
    if torch.cuda.is_available():
        gradchecking(use_cuda=True)

================================================
FILE: install.sh
================================================
python setup.py build develop

# or if you are on macOS
# MACOSX_DEPLOYMENT_TARGET=10.9 CC=clang CXX=clang++ python setup.py build develop

================================================
FILE: setup.py
================================================
#!/usr/bin/env python

import glob
import os

import torch
from setuptools import find_packages
from setuptools import setup
from torch.utils.cpp_extension import CUDA_HOME
from torch.utils.cpp_extension import CppExtension
from torch.utils.cpp_extension import CUDAExtension

requirements = ["torch", "torchvision"]


def get_extensions():
    this_dir = os.path.dirname(os.path.abspath(__file__))
    extensions_dir = os.path.join(this_dir, "gcn", "csrc")

    main_file = glob.glob(os.path.join(extensions_dir, "*.cpp"))
    source_cpu = glob.glob(os.path.join(extensions_dir, "cpu", "*.cpp"))
    source_cuda = glob.glob(os.path.join(extensions_dir, "cuda", "*.cu"))

    sources = main_file + source_cpu
    extension = CppExtension

    extra_compile_args = {"cxx": []}
    define_macros = []

    if torch.cuda.is_available() and CUDA_HOME is not None:
        extension = CUDAExtension
        sources += source_cuda
        define_macros += [("WITH_CUDA", None)]
        extra_compile_args["nvcc"] = [
            "-DCUDA_HAS_FP16=1",
            "-D__CUDA_NO_HALF_OPERATORS__",
            "-D__CUDA_NO_HALF_CONVERSIONS__",
            "-D__CUDA_NO_HALF2_OPERATORS__",
        ]

    sources = [os.path.join(extensions_dir, s) for s in sources]

    include_dirs = [extensions_dir]

    ext_modules = [
        extension(
            "gcn._C",
            sources,
            include_dirs=include_dirs,
            define_macros=define_macros,
            extra_compile_args=extra_compile_args,
        )
    ]

    return ext_modules


setup(
    name="gcn",
    version="0.1",
    author="gujiaxin",
    url="https://github.com/jxgu1016/Gabor_CNN_PyTorch",
    description="Gabor Convolutional Networks in pytorch",
    packages=find_packages(),
    # install_requires=requirements,
    ext_modules=get_extensions(),
    cmdclass={"build_ext": torch.utils.cpp_extension.BuildExtension},
)