def create_layers():
layers = []
input_dim = input_size
for i, hidden_size in enumerate(hidden_sizes):
fc = nn.Conv1d(
in_channels=input_dim * num_heads,
out_channels=hidden_size * num_heads,
kernel_size=1,
groups=num_heads,
)
layers.append(fc)
if isinstance(hidden_activation, str):
activation = activation_from_string(hidden_activation)
else:
activation = hidden_activation
layers.append(activation)
input_dim = hidden_size
last_fc = nn.Conv1d(
in_channels=input_dim * num_heads,
out_channels=output_size_per_mlp * num_heads,
kernel_size=1,
groups=num_heads,
)
layers.append(last_fc)
if output_activation != "identity":
if isinstance(output_activation, str):
activation = activation_from_string(output_activation)
else:
activation = output_activation
layers.append(activation)
return layers
def __init__(
self,
output_sizes,
output_activations=None,
):
super().__init__()
if output_activations is None:
output_activations = ['identity' for _ in output_sizes]
else:
if len(output_activations) != len(output_sizes):
raise ValueError("output_activation and output_sizes must have "
"the same length")
self._output_narrow_params = []
self._output_activations = []
for output_activation in output_activations:
if isinstance(output_activation, str):
output_activation = activation_from_string(output_activation)
self._output_activations.append(output_activation)
start_idx = 0
for output_size in output_sizes:
self._output_narrow_params.append((start_idx, output_size))
start_idx = start_idx + output_size
def __init__(
self,
input_width,
input_height,
input_channels,
kernel_sizes,
n_channels,
strides,
paddings,
normalization_type="none",
hidden_init=None,
hidden_activation="relu",
output_activation=identity,
pool_type="none",
pool_sizes=None,
pool_strides=None,
pool_paddings=None,
):
assert len(kernel_sizes) == len(n_channels) == len(strides) == len(
paddings)
assert normalization_type in {"none", "batch", "layer"}
assert pool_type in {"none", "max2d"}
if pool_type == "max2d":
assert len(pool_sizes) == len(pool_strides) == len(pool_paddings)
super().__init__()
self.output_activation = output_activation
if isinstance(hidden_activation, str):
hidden_activation = activation_from_string(hidden_activation)
self.hidden_activation = hidden_activation
self.input_width = input_width
self.input_height = input_height
self.input_channels = input_channels
self.normalization_type = normalization_type
self.layers = nn.ModuleList()
self.pool_layers = nn.ModuleList()
self.norm_layers = nn.ModuleList()
self.pool_type = pool_type
for i, (out_channels, kernel_size, stride, padding) in enumerate(
zip(n_channels, kernel_sizes, strides, paddings)):
deconv = nn.ConvTranspose2d(
input_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding,
)
if hidden_init:
hidden_init(deconv.weight)
layer = deconv
self.layers.append(layer)
input_channels = out_channels
if pool_type == "max2d":
if pool_sizes[i] > 1:
self.pool_layers.append(
nn.MaxUnpool2d(
kernel_size=pool_sizes[i],
stride=pool_strides[i],
padding=pool_paddings[i],
))
else:
self.pool_layers.append(None)
test_mat = torch.zeros(
1,
self.input_channels,
self.input_width,
self.input_height,
)
for layer in self.layers:
test_mat = layer(test_mat)
if self.normalization_type == "batch":
self.norm_layers.append(nn.BatchNorm2d(test_mat.shape[1]))
if self.normalization_type == "layer":
self.norm_layers.append(nn.LayerNorm(test_mat.shape[1:]))
self.output_shape = test_mat.shape[1:] # ignore batch dim
def __init__(
self,
input_width,
input_height,
input_channels,
kernel_sizes,
n_channels,
strides,
paddings,
normalization_type='none',
hidden_init=None,
hidden_activation='relu',
output_activation=identity,
pool_type='none',
pool_sizes=None,
pool_strides=None,
pool_paddings=None,
):
assert len(kernel_sizes) == \
len(n_channels) == \
len(strides) == \
len(paddings)
assert normalization_type in {'none', 'batch', 'layer'}
assert pool_type in {'none', 'max2d'}
if pool_type == 'max2d':
assert len(pool_sizes) == len(pool_strides) == len(pool_paddings)
super().__init__()
self.input_width = input_width
self.input_height = input_height
self.input_channels = input_channels
self.output_activation = output_activation
if isinstance(hidden_activation, str):
hidden_activation = activation_from_string(hidden_activation)
self.hidden_activation = hidden_activation
self.normalization_type = normalization_type
self.conv_input_length = self.input_width * self.input_height * self.input_channels
self.pool_type = pool_type
self.conv_layers = nn.ModuleList()
self.conv_norm_layers = nn.ModuleList()
self.pool_layers = nn.ModuleList()
for i, (out_channels, kernel_size, stride, padding) in enumerate(
zip(n_channels, kernel_sizes, strides, paddings)):
conv = nn.Conv2d(input_channels,
out_channels,
kernel_size,
stride=stride,
padding=padding)
if hidden_init:
hidden_init(conv.weight)
conv_layer = conv
self.conv_layers.append(conv_layer)
input_channels = out_channels
if pool_type == 'max2d':
if pool_sizes[i] > 1:
self.pool_layers.append(
nn.MaxPool2d(
kernel_size=pool_sizes[i],
stride=pool_strides[i],
padding=pool_paddings[i],
))
else:
self.pool_layers.append(None)
# use torch rather than ptu because initially the model is on CPU
test_mat = torch.zeros(
1,
self.input_channels,
self.input_height,
self.input_width,
)
# find output dim of conv_layers by trial and add norm conv layers
for i, conv_layer in enumerate(self.conv_layers):
test_mat = conv_layer(test_mat)
if self.normalization_type == 'batch':
self.conv_norm_layers.append(nn.BatchNorm2d(test_mat.shape[1]))
if self.normalization_type == 'layer':
self.conv_norm_layers.append(nn.LayerNorm(test_mat.shape[1:]))
if self.pool_type != 'none':
if self.pool_layers[i]:
test_mat = self.pool_layers[i](test_mat)
self.output_shape = test_mat.shape[1:] # ignore batch dim