Spaces:
Runtime error
Runtime error
File size: 3,119 Bytes
01bb3bb |
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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import torch.nn as nn
import torch
class CircConv(nn.Module):
def __init__(self, state_dim, out_state_dim=None, n_adj=4):
super(CircConv, self).__init__()
self.n_adj = n_adj
out_state_dim = state_dim if out_state_dim is None else out_state_dim
self.fc = nn.Conv1d(state_dim, out_state_dim, kernel_size=self.n_adj*2+1)
def forward(self, input, adj):
input = torch.cat([input[..., -self.n_adj:], input, input[..., :self.n_adj]], dim=2)
return self.fc(input)
class DilatedCircConv(nn.Module):
def __init__(self, state_dim, out_state_dim=None, n_adj=4, dilation=1):
super(DilatedCircConv, self).__init__()
self.n_adj = n_adj
self.dilation = dilation
out_state_dim = state_dim if out_state_dim is None else out_state_dim
self.fc = nn.Conv1d(state_dim, out_state_dim, kernel_size=self.n_adj*2+1, dilation=self.dilation)
def forward(self, input, adj):
if self.n_adj != 0:
input = torch.cat([input[..., -self.n_adj*self.dilation:], input, input[..., :self.n_adj*self.dilation]], dim=2)
return self.fc(input)
_conv_factory = {
'grid': CircConv,
'dgrid': DilatedCircConv
}
class BasicBlock(nn.Module):
def __init__(self, state_dim, out_state_dim, conv_type, n_adj=4, dilation=1):
super(BasicBlock, self).__init__()
self.conv = _conv_factory[conv_type](state_dim, out_state_dim, n_adj, dilation)
self.relu = nn.ReLU(inplace=True)
self.norm = nn.BatchNorm1d(out_state_dim)
def forward(self, x, adj=None):
x = self.conv(x, adj)
x = self.relu(x)
x = self.norm(x)
return x
class DeepSnake(nn.Module):
def __init__(self, state_dim, feature_dim, conv_type='dgrid'):
super(DeepSnake, self).__init__()
self.head = BasicBlock(feature_dim, state_dim, conv_type)
self.res_layer_num = 7
dilation = [1, 1, 1, 2, 2, 4, 4]
for i in range(self.res_layer_num):
conv = BasicBlock(state_dim, state_dim, conv_type, n_adj=4, dilation=dilation[i])
self.__setattr__('res'+str(i), conv)
fusion_state_dim = 256
self.fusion = nn.Conv1d(state_dim * (self.res_layer_num + 1), fusion_state_dim, 1)
self.prediction = nn.Sequential(
nn.Conv1d(state_dim * (self.res_layer_num + 1) + fusion_state_dim, 256, 1),
nn.ReLU(inplace=True),
nn.Conv1d(256, 64, 1),
nn.ReLU(inplace=True),
nn.Conv1d(64, 2, 1)
)
def forward(self, x, adj):
states = []
x = self.head(x, adj)
states.append(x)
for i in range(self.res_layer_num):
x = self.__getattr__('res'+str(i))(x, adj) + x
states.append(x)
state = torch.cat(states, dim=1)
global_state = torch.max(self.fusion(state), dim=2, keepdim=True)[0]
global_state = global_state.expand(global_state.size(0), global_state.size(1), state.size(2))
state = torch.cat([global_state, state], dim=1)
x = self.prediction(state)
return x
|