SOC-YOLO/utils/softpool.py at main · Human-Language-Intelligence/SOC-YOLO · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
from torch import nn
from torch.autograd import Function
import torch.nn.functional as F
import torch
from torch.nn.modules.utils import _triple, _pair, _single
from torch.utils.cpp_extension import load

#import softpool_cuda
class SoftPool2d(torch.nn.Module):
    def __init__(self,kernel_size,stride=None,padding=0,ceil_mode = False,count_include_pad = True,divisor_override = None):
        super(SoftPool2d, self).__init__()
        self.avgpool = torch.nn.AvgPool2d(kernel_size,stride,padding,ceil_mode,count_include_pad,divisor_override)
    def forward(self, x):
        x_exp = torch.exp(x)
        x_exp_pool = self.avgpool(x_exp)
        x = self.avgpool(x_exp*x)
        return x/x_exp_pool

class SoftPooling2D(torch.nn.Module):
    def __init__(self,kernel_size,stride=1,padding=5//2,ceil_mode = False,count_include_pad = True,divisor_override = None):
        super(SoftPooling2D, self).__init__()
        self.avgpool = torch.nn.AvgPool2d(kernel_size,stride,padding,ceil_mode,count_include_pad,divisor_override)

    def forward(self, x):

        x_exp = torch.exp(x)
        x_exp_pool = self.avgpool(x_exp)
        x = self.avgpool(x_exp*x)
        return x/x_exp_pool