|
|
|
|
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
class MLP(nn.Module): |
|
def __init__(self, |
|
filter_channels, |
|
merge_layer=0, |
|
res_layers=[], |
|
norm='group', |
|
last_op=None): |
|
super(MLP, self).__init__() |
|
|
|
self.filters = nn.ModuleList() |
|
self.norms = nn.ModuleList() |
|
self.merge_layer = merge_layer if merge_layer > 0 else len(filter_channels) // 2 |
|
self.res_layers = res_layers |
|
self.norm = norm |
|
self.last_op = last_op |
|
|
|
for l in range(0, len(filter_channels)-1): |
|
if l in self.res_layers: |
|
self.filters.append(nn.Conv1d( |
|
filter_channels[l] + filter_channels[0], |
|
filter_channels[l+1], |
|
1)) |
|
else: |
|
self.filters.append(nn.Conv1d( |
|
filter_channels[l], |
|
filter_channels[l+1], |
|
1)) |
|
if l != len(filter_channels)-2: |
|
if norm == 'group': |
|
self.norms.append(nn.GroupNorm(32, filter_channels[l+1])) |
|
elif norm == 'batch': |
|
self.norms.append(nn.BatchNorm1d(filter_channels[l+1])) |
|
|
|
def forward(self, feature): |
|
''' |
|
feature may include multiple view inputs |
|
args: |
|
feature: [B, C_in, N] |
|
return: |
|
[B, C_out, N] prediction |
|
''' |
|
y = feature |
|
tmpy = feature |
|
phi = None |
|
for i, f in enumerate(self.filters): |
|
y = f( |
|
y if i not in self.res_layers |
|
else torch.cat([y, tmpy], 1) |
|
) |
|
if i != len(self.filters)-1: |
|
if self.norm not in ['batch', 'group']: |
|
y = F.leaky_relu(y) |
|
else: |
|
y = F.leaky_relu(self.norms[i](y)) |
|
if i == self.merge_layer: |
|
phi = y.clone() |
|
|
|
if self.last_op is not None: |
|
y = self.last_op(y) |
|
|
|
return y, phi |
|
|
