81 lines
3.0 KiB
Python
81 lines
3.0 KiB
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import numpy as np
|
||
|
|
import torch.nn.functional as F
|
||
|
|
|
||
|
|
def conv3x3(in_planes, out_planes, stride=1):
|
||
|
|
"""3x3 convolution with padding"""
|
||
|
|
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||
|
|
padding=1, bias=False)
|
||
|
|
|
||
|
|
def conv5x5(in_planes, out_planes, stride=1):
|
||
|
|
"""5x5 convolution with padding"""
|
||
|
|
return nn.Conv2d(in_planes, out_planes, kernel_size=5, stride=stride,
|
||
|
|
padding=2, bias=False)
|
||
|
|
|
||
|
|
def conv7x7(in_planes, out_planes, stride=1):
|
||
|
|
"""7x7 convolution with padding"""
|
||
|
|
return nn.Conv2d(in_planes, out_planes, kernel_size=7, stride=stride,
|
||
|
|
padding=3, bias=False)
|
||
|
|
|
||
|
|
def conv1x1(in_planes, out_planes, stride=1):
|
||
|
|
"""1x1 convolution"""
|
||
|
|
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
|
||
|
|
|
||
|
|
class SequenceDistributed(nn.Module):
|
||
|
|
def __init__(self, module, batch_first=False):
|
||
|
|
super(SequenceDistributed, self).__init__()
|
||
|
|
self.module = module
|
||
|
|
self.batch_first = batch_first
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
if len(x.size()) <= 2:
|
||
|
|
return self.module(x)
|
||
|
|
|
||
|
|
# Squash samples and timesteps into a single axis
|
||
|
|
x_reshape = x.contiguous().view(-1, x.size(-1)) # (samples * timesteps, input_size)
|
||
|
|
|
||
|
|
y = self.module(x_reshape)
|
||
|
|
|
||
|
|
# We have to reshape Y
|
||
|
|
if self.batch_first:
|
||
|
|
y = y.contiguous().view(x.size(0), -1, y.size(-1)) # (samples, timesteps, output_size)
|
||
|
|
else:
|
||
|
|
y = y.view(-1, x.size(1), y.size(-1)) # (timesteps, samples, output_size)
|
||
|
|
|
||
|
|
return y
|
||
|
|
|
||
|
|
# Input into this block is of shape (sequence, filters, width, height)
|
||
|
|
# Output is (attention_hidden_size, width, height)
|
||
|
|
class ConvAttentionBlock(nn.Module):
|
||
|
|
|
||
|
|
def __init__(self, planes, attention_hidden_size=8, query_conv=conv1x1, key_conv=conv1x1, value_conv=conv1x1):
|
||
|
|
super(ConvAttentionBlock, self).__init__()
|
||
|
|
self.query_conv_dist = SequenceDistributed(query_conv(planes, attention_hidden_size))
|
||
|
|
self.key_conv_dist = SequenceDistributed(key_conv(planes, attention_hidden_size))
|
||
|
|
self.value_conv_dist = value_conv(planes, attention_hidden_size)
|
||
|
|
self.hidden_size = attention_hidden_size
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
# All values come out of this with the shape (batch, sequence, hidden, width, height)
|
||
|
|
query = self.query_conv_dist(x)
|
||
|
|
key = self.key_conv_dist(x)
|
||
|
|
value = self.value_conv_dist(x)
|
||
|
|
|
||
|
|
# Permute to (batch, width, height, sequence, hidden)
|
||
|
|
query = query.permute(0, 3, 4, 1, 2)
|
||
|
|
key = key.permute(0, 3, 4, 1, 2)
|
||
|
|
value = value.permute(0, 3, 4, 1, 2)
|
||
|
|
|
||
|
|
# Perform attention operation.
|
||
|
|
scores = torch.matmul(query, key.transpose(-2, -1)) / torch.sqrt(self.hidden_size)
|
||
|
|
scores = torch.softmax(scores, dim=-1)
|
||
|
|
result = torch.matmul(scores, value)
|
||
|
|
|
||
|
|
# Collapse out the sequence dim.
|
||
|
|
result = torch.sum(result, dim=-2)
|
||
|
|
|
||
|
|
# Permute back to (batch, hidden, width, height)
|
||
|
|
result = result.permute(0, 3, 1, 2)
|
||
|
|
return result
|