From 27c7334e8176d1b0de0aab9e8616c48ba21903f8 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 2 Apr 2020 12:22:15 -0700 Subject: [PATCH] new layers.py file --- models.py | 67 +++---------------------------------------------- utils/layers.py | 62 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 65 insertions(+), 64 deletions(-) create mode 100644 utils/layers.py diff --git a/models.py b/models.py index fee58077..089b6034 100755 --- a/models.py +++ b/models.py @@ -1,8 +1,6 @@ -import torch.nn.functional as F - from utils.google_utils import * +from utils.layers import * from utils.parse_config import * -from utils.utils import * ONNX_EXPORT = False @@ -70,7 +68,7 @@ def create_modules(module_defs, img_size): layers = mdef['from'] filters = output_filters[-1] routs.extend([i + l if l < 0 else l for l in layers]) - modules = weightedFeatureFusion(layers=layers, weight='weights_type' in mdef) + modules = WeightedFeatureFusion(layers=layers, weight='weights_type' in mdef) elif mdef['type'] == 'reorg3d': # yolov3-spp-pan-scale pass @@ -111,65 +109,6 @@ def create_modules(module_defs, img_size): return module_list, routs_binary -class weightedFeatureFusion(nn.Module): # weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 - def __init__(self, layers, weight=False): - super(weightedFeatureFusion, self).__init__() - self.layers = layers # layer indices - self.weight = weight # apply weights boolean - self.n = len(layers) + 1 # number of layers - if weight: - self.w = torch.nn.Parameter(torch.zeros(self.n)) # layer weights - - def forward(self, x, outputs): - # Weights - if self.weight: - w = torch.sigmoid(self.w) * (2 / self.n) # sigmoid weights (0-1) - x = x * w[0] - - # Fusion - nx = x.shape[1] # input channels - for i in range(self.n - 1): - a = outputs[self.layers[i]] * w[i + 1] if self.weight else outputs[self.layers[i]] # feature to add - na = a.shape[1] # feature channels - - # Adjust channels - if nx == na: # same shape - x = x + a - elif nx > na: # slice input - x[:, :na] = x[:, :na] + a # or a = nn.ZeroPad2d((0, 0, 0, 0, 0, dc))(a); x = x + a - else: # slice feature - x = x + a[:, :nx] - - return x - - -class SwishImplementation(torch.autograd.Function): - @staticmethod - def forward(ctx, i): - ctx.save_for_backward(i) - return i * torch.sigmoid(i) - - @staticmethod - def backward(ctx, grad_output): - sigmoid_i = torch.sigmoid(ctx.saved_variables[0]) - return grad_output * (sigmoid_i * (1 + ctx.saved_variables[0] * (1 - sigmoid_i))) - - -class MemoryEfficientSwish(nn.Module): - def forward(self, x): - return SwishImplementation.apply(x) - - -class Swish(nn.Module): - def forward(self, x): - return x.mul_(torch.sigmoid(x)) - - -class Mish(nn.Module): # https://github.com/digantamisra98/Mish - def forward(self, x): - return x.mul_(F.softplus(x).tanh()) - - class YOLOLayer(nn.Module): def __init__(self, anchors, nc, img_size, yolo_index, layers): super(YOLOLayer, self).__init__() @@ -277,7 +216,7 @@ class Darknet(nn.Module): l = [i - 1] + module.layers # layers s = [list(x.shape)] + [list(out[i].shape) for i in module.layers] # shapes str = ' >> ' + ' + '.join(['layer %g %s' % x for x in zip(l, s)]) - x = module(x, out) # weightedFeatureFusion() + x = module(x, out) # WeightedFeatureFusion() elif mtype == 'route': # concat layers = mdef['layers'] if verbose: diff --git a/utils/layers.py b/utils/layers.py new file mode 100644 index 00000000..cb72d67b --- /dev/null +++ b/utils/layers.py @@ -0,0 +1,62 @@ +import torch.nn.functional as F + +from utils.utils import * + + +class WeightedFeatureFusion(nn.Module): # weighted sum of 2 or more layers https://arxiv.org/abs/1911.09070 + def __init__(self, layers, weight=False): + super(WeightedFeatureFusion, self).__init__() + self.layers = layers # layer indices + self.weight = weight # apply weights boolean + self.n = len(layers) + 1 # number of layers + if weight: + self.w = torch.nn.Parameter(torch.zeros(self.n), requires_grad=True) # layer weights + + def forward(self, x, outputs): + # Weights + if self.weight: + w = torch.sigmoid(self.w) * (2 / self.n) # sigmoid weights (0-1) + x = x * w[0] + + # Fusion + nx = x.shape[1] # input channels + for i in range(self.n - 1): + a = outputs[self.layers[i]] * w[i + 1] if self.weight else outputs[self.layers[i]] # feature to add + na = a.shape[1] # feature channels + + # Adjust channels + if nx == na: # same shape + x = x + a + elif nx > na: # slice input + x[:, :na] = x[:, :na] + a # or a = nn.ZeroPad2d((0, 0, 0, 0, 0, dc))(a); x = x + a + else: # slice feature + x = x + a[:, :nx] + + return x + + +class SwishImplementation(torch.autograd.Function): + @staticmethod + def forward(ctx, i): + ctx.save_for_backward(i) + return i * torch.sigmoid(i) + + @staticmethod + def backward(ctx, grad_output): + sigmoid_i = torch.sigmoid(ctx.saved_variables[0]) + return grad_output * (sigmoid_i * (1 + ctx.saved_variables[0] * (1 - sigmoid_i))) + + +class MemoryEfficientSwish(nn.Module): + def forward(self, x): + return SwishImplementation.apply(x) + + +class Swish(nn.Module): + def forward(self, x): + return x.mul_(torch.sigmoid(x)) + + +class Mish(nn.Module): # https://github.com/digantamisra98/Mish + def forward(self, x): + return x.mul_(F.softplus(x).tanh())