def __init__(self, output_size, filter_size=5):
super(network, self).__init__()
features = [(1,),
(4, 4, 4, 4),
(4, 4, 4, 4),
(4, 4, 4, 4),
(output_size,)]
common_block_params = {
'size': filter_size,
'padding': filter_size//2,
'stride': 1,
'normalization': 'instance',
'radial_window': partial(
basis_kernels.gaussian_window_fct_convenience_wrapper,
mode='compromise', border_dist=0, sigma=0.6),
}
block_params = [
{'activation': (F.relu, F.sigmoid)},
{'activation': (F.relu, F.sigmoid)},
{'activation': (F.relu, F.sigmoid)},
{'activation': None},
]
assert len(block_params) + 1 == len(features)
blocks = [GatedBlock(features[i], features[i + 1],
**common_block_params, **block_params[i])
for i in range(len(block_params))]
self.layers = torch.nn.Sequential(
*blocks,
)
def __init__(self):
super().__init__()
features = [(1, ), (4, 2, 1), (4, 2, 1), (2, )]
block_params = [
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': None
},
]
common_block_params = {
'size': 7,
'stride': 1,
'padding': 9,
'dilation': 3,
'normalization': None,
}
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = torch.nn.Sequential(*blocks)
self.post_activations = None
def __init__(self):
super().__init__()
# The parameters of a GatedBlock are:
# - The representation multiplicities (scalar, vector and dim. 5 repr.) for the input and the output
# - the non linearities for the scalars and the gates (None for no non-linearity)
# - stride, padding... same as 2D convolution
features = [
(1, ), # As input we have a scalar field
(
2, 2, 2, 2
), # Note that this particular choice of multiplicities it completely arbitrary
(4, 4, 3, 3),
(4, 4, 3, 3),
(4, 4, 3, 3),
(4, 4, 3, 3),
(4, 4, 3, 3),
(20, ) # scalar fields to end with fully-connected layers
]
common_block_params = {
'size': 5,
'stride': 2,
'padding': 3,
}
block_params = [
{
'activation': (None, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': None
},
]
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = nn.Sequential(*blocks, AvgSpacial(), nn.Linear(20, 50),
nn.ReLU(), nn.Linear(50, 2),
nn.LogSoftmax(dim=1))
def __init__(self, smooth_stride):
super().__init__()
features = [(1, ), (4, 4, 4, 1), (16, 16, 16, 0), (32, 16, 16, 0),
(128, )]
common_block_params = {
'size': 5,
'padding': 4,
'activation': (F.relu, torch.sigmoid),
'smooth_stride': smooth_stride,
}
block_params = [
{
'stride': 1
},
{
'stride': 2
},
{
'stride': 2
},
{
'stride': 1
},
]
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(features) - 1)
]
self.sequence = torch.nn.Sequential(*blocks, AvgSpacial(),
nn.Dropout(p=.2),
nn.Linear(features[-1][0], 8))
def __init__(self):
super(SE3Net, self).__init__()
features = [(1, ), (2, 2, 2, 2), (4, 4, 4, 4), (16, )]
common_block_params = {
'size': 3,
'stride': 1,
'padding': 1,
'capsule_dropout_p': 0.01,
'normalization': 'batch',
}
block_params = [
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
]
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = torch.nn.Sequential(*blocks, AvgSpacial(),
torch.nn.Linear(16, 10))
def __init__(self):
super().__init__()
features = [
(1, ),
(2, 2, 2, 2),
(0, 2, 2, 2), # no scalar fields here
(2, 2, 2, 2),
(1, ),
]
common_block_params = {
'size': 5,
'padding': 2,
}
block_params = [
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': None
},
]
assert len(block_params) + 1 == len(features)
self.blocks = nn.Sequential(*(GatedBlock(features[i], features[
i + 1], **common_block_params, **block_params[i])
for i in range(len(block_params))))
def __init__(self, n_in, n_out):
super().__init__()
self.int_repr = None
features = [
(n_in, ), # 77
(12, 4, 1),
(24, 8, 2),
(24, 8, 2),
(24, 8, 2),
(n_out, ),
]
common_block_params = {
'size': 7,
'padding': 3,
'normalization': 'batch_max',
'smooth_stride': True,
'smooth_stride': False,
}
block_params = [
{
'activation': (F.relu, torch.sigmoid),
'stride': 2
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': None
},
]
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
for p in blocks[-1].parameters():
nn.init.zeros_(p)
self.sequence = torch.nn.Sequential(
*blocks,
AvgSpacial(),
)
def __init__(self, features, stride=1):
super().__init__()
n = len(features) - 1
self.sequential = nn.Sequential(*[
GatedBlock(features[i],
features[i + 1],
size=7,
padding=3,
stride=stride if i == 0 else 1,
n_radial=2,
activation=F.relu) for i in range(n)
])
self.shortcut = GatedBlock(features[0],
features[-1],
size=1,
stride=stride,
n_radial=1,
activation=None)
def test_grad(self):
m = GatedBlock([1, 1, 1], [1, 1, 1],
5,
activation=(torch.relu, torch.sigmoid),
checkpoint=False).type(torch.float64)
x = torch.rand(1,
1 + 3 + 5,
6,
6,
6,
requires_grad=True,
dtype=torch.float64)
self.assertTrue(torch.autograd.gradcheck(m, (x, ), eps=1e-3))
def __init__(self, n_out):
super().__init__()
self.int_repr = None
features = [ # (6) double all channels
(1, ), (4, 4, 4), (8, 8, 8), (16, 16, 16), (32, 16, 16), (200, )
]
common_block_params = {
'size': 5, # (5) = 5->7
'stride': 2,
'padding': 3,
'normalization': 'batch',
'capsule_dropout_p': 0.1 # (4) Maurice suggestion
}
block_params = [
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
]
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = torch.nn.Sequential(
*blocks,
AvgSpacial(),
nn.Linear(features[-1][0], n_out),
)
def __init__(self):
super().__init__()
features = [(1, ), (4, 4, 4), (8, 8, 8), (16, 16, 16), (200, )]
radial_window = partial(
basis_kernels.gaussian_window_fct_convenience_wrapper,
mode='compromise',
border_dist=0,
sigma=0.6)
common_block_params = {
'size': 5,
'stride': 2,
'padding': 3,
'normalization': 'batch',
'radial_window': radial_window,
'capsule_dropout_p': 0.1
}
block_params = [
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
{
'activation': (F.relu, F.sigmoid)
},
]
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = torch.nn.Sequential(*blocks, AvgSpacial(),
nn.Linear(features[-1][0], 10),
nn.LogSoftmax(dim=1))
def __init__(self):
super(SE3Net, self).__init__()
features = [(1, ), (2, 2, 2, 1), (4, 4, 4, 0), (6, 4, 4, 0), (64, )]
common_block_params = {
'size': 5,
'padding': 4,
'dilation': 2,
'activation': (F.relu, torch.sigmoid),
'capsule_dropout_p': .0
}
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params)
for i in range(len(features) - 1)
]
self.sequence = torch.nn.Sequential(*blocks, AvgSpacial(),
nn.Dropout(p=.2),
nn.Linear(64, 10))
def __init__(self):
super().__init__()
features = [(1, 2), (4, 4, 4), (8, 8, 4), (16, 16, 4), (32, 16, 4),
(0, 2)]
common_block_params = {
'size': 5,
'stride': 2,
'padding': 4,
'normalization': None,
'capsule_dropout_p': 0.0,
'smooth_stride': False,
}
block_params = [
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid)
},
]
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = torch.nn.Sequential(
*blocks,
AvgSpacial(),
)
def __init__(self):
super(SE3Net, self).__init__()
# features = [
# (1,),
# (2, 2, 2, 1),
# (4, 4, 4, 0),
# (6, 4, 4, 0),
# (64,)
# ]
# common_block_params = {
# 'size': 5,
# 'padding': 4,
# 'dilation': 2,
# 'activation': (F.relu, torch.sigmoid),
# }
features = [(1, ), (2, 2, 2, 1), (8, 8, 8, 0), (16, 8, 8, 0), (64, )]
common_block_params = {
'size': 5,
'padding': 4,
'activation': (F.relu, torch.sigmoid),
'smooth_stride': True,
}
block_params = [
{
'stride': 1
},
{
'stride': 2
},
{
'stride': 2
},
{
'stride': 1
},
]
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(features) - 1)
]
self.sequence = torch.nn.Sequential(*blocks, AvgSpacial(),
nn.Dropout(p=.2),
nn.Linear(64, 10))
def ResNet():
return nn.Sequential(
GatedBlock((1, ), (5, 3, 1), size=7, n_radial=2, activation=F.relu),
ResNetBlock([(5, 3, 1), (1, 1, 1), (5, 3, 1)], stride=2),
ResNetBlock([(5, 3, 1), (1, 1, 1), (5, 3, 1)], stride=2),
ResNetBlock([(5, 3, 1), (1, 1, 1), (2, )], stride=2), AvgSpacial())
def __init__(self, output_size, filter_size=5):
super(network, self).__init__()
size = filter_size
common_params = {
'radial_window':
partial(basis_kernels.gaussian_window_fct_convenience_wrapper,
mode='compromise',
border_dist=0,
sigma=0.6),
'batch_norm_momentum':
0.01,
}
features = [
(1, ),
# (12, 12, 12),
# (24, 24, 24),
# (48, 48, 48),
# (24, 24, 24),
# (12, 12, 12),
(8, 8, 8, 4),
(16, 16, 16, 8),
(32, 32, 32, 16),
(16, 16, 16, 8),
(8, 8, 8, 4),
(output_size, )
]
# TODO: do padding using ReplicationPad3d?
# TODO: on validation - use overlapping patches and only use center of patch
self.conv1 = nn.Sequential(
GatedBlock(features[0],
features[1],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params),
GatedBlock(features[1],
features[1],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.conv2 = nn.Sequential(
GatedBlock(features[1],
features[2],
size=size,
padding=size // 2,
stride=2,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params),
GatedBlock(features[2],
features[2],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.conv3 = nn.Sequential(
GatedBlock(features[2],
features[3],
size=size,
padding=size // 2,
stride=2,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params),
GatedBlock(features[3],
features[3],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.up1 = nn.Sequential(
nn.Upsample(scale_factor=2, mode="nearest"),
GatedBlock(features[3],
features[4],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.merge1 = Merge()
self.conv4 = nn.Sequential(
GatedBlock(features[3],
features[4],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params),
GatedBlock(features[4],
features[4],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.up2 = nn.Sequential(
nn.Upsample(scale_factor=2, mode="nearest"),
GatedBlock(features[4],
features[5],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.merge2 = Merge()
self.conv5 = nn.Sequential(
GatedBlock(features[4],
features[5],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params),
GatedBlock(features[5],
features[5],
size=size,
padding=size // 2,
stride=1,
activation=(F.relu, F.sigmoid),
normalization="instance",
**common_params))
self.conv_final = GatedBlock(features[5],
features[6],
size=1,
padding=0,
stride=1,
activation=None,
normalization=None,
**common_params)
def __init__(self, config, nfeat):
"""
Build the netwrok in a flexible way based on `config` string, which contains sequence
of comma-delimited layer definiton tokens layer_arg1_arg2_... See README.md for examples
:param config:
:param nfeat: Number of input channels
"""
super().__init__()
self.register_buffer('scaler_mean', torch.zeros(1, nfeat, 1, 1, 1))
self.register_buffer('scaler_std', torch.ones(1, nfeat, 1, 1, 1))
for d, conf in enumerate(config.split(',')):
conf = conf.strip().split('_')
if conf[0] == 'b': # Batch norm;
self.add_module(str(d), nn.BatchNorm3d(nfeat))
elif conf[0] == 'r': # ReLU
self.add_module(str(d), nn.ReLU(True))
elif conf[0] == 'm': # Max pooling
kernel_size = int(conf[1])
self.add_module(str(d), nn.MaxPool3d(kernel_size))
elif conf[0] == 'a': # Avg pooling
kernel_size = int(conf[1])
self.add_module(str(d), nn.AvgPool3d(kernel_size))
elif conf[
0] == 'c': # 3D convolution args: output feat, kernel size, padding, stride
nfeato = int(conf[1])
kernel_size = int(conf[2])
padding = int(conf[3]) if len(conf) > 3 else 0
stride = int(conf[4]) if len(conf) > 4 else 1
self.add_module(
str(d),
nn.Conv3d(nfeat, nfeato, kernel_size, stride, padding))
nfeat = nfeato
elif conf[
0] == 'se': # SE(3)-covariant block args: output feat, mult1, mult2, mult3, kernel size, padding, stride, bnnorm, smoothing
nfeato = int(conf[1])
mult1 = int(conf[2])
mult2 = int(conf[3])
mult3 = int(conf[4])
kernel_size = int(conf[5])
padding = int(conf[6]) if len(conf) > 6 else 0
stride = int(conf[7]) if len(conf) > 7 else 1
normalization = conf[8] if len(conf) > 8 else None
smooth = bool(int(conf[9])) if len(conf) > 9 else False
if isinstance(nfeat, int):
nfeat = (nfeat, )
nfeato = tuple(
[n for n in (nfeato, mult1, mult2, mult3) if n > 0])
activation = (None, nnf.sigmoid)
elif mult1 <= 0:
nfeato = (nfeato, )
activation = None
else:
nfeato = tuple(
[n for n in (nfeato, mult1, mult2, mult3) if n > 0])
activation = (nnf.relu, nnf.sigmoid)
conv = GatedBlock(nfeat,
nfeato,
size=kernel_size,
padding=padding,
stride=stride,
activation=activation,
normalization=normalization,
smooth_stride=smooth)
self.add_module(str(d), conv)
if mult1 <= 0:
nfeato = nfeato[0]
nfeat = nfeato
else:
raise NotImplementedError('Unknown module: ' + conf[0])
self.nfeato = nfeat
def __init__(self, in_repr, out_reprs,
size=3,
stride=1,
radial_window=None,
batch_norm_momentum=0.01,
normalization="batch",
capsule_dropout_p=0.1,
scalar_gate_activation=(F.relu, F.sigmoid),
downsample_by_pooling=False):
super().__init__()
reprs = [in_repr] + out_reprs
self.layers = []
single_layer = len(out_reprs) == 1
conv_stride = 1 if downsample_by_pooling else stride
for i in range(len(reprs) - 1):
# No activation in last block
activation = scalar_gate_activation
if i == (len(reprs) - 2) and not single_layer:
activation = None
self.layers.append(
GatedBlock(reprs[i], reprs[i + 1],
size=size, padding=size//2,
stride=conv_stride if i == 0 else 1,
activation=activation,
radial_window=radial_window,
batch_norm_momentum=batch_norm_momentum,
normalization=normalization,
capsule_dropout_p=capsule_dropout_p))
if downsample_by_pooling and i == 0 and stride > 1:
self.layers.append(nn.AvgPool3d(kernel_size=size,
padding=size//2,
stride=stride))
self.layers = nn.Sequential(*self.layers)
self.shortcut = None
self.activation = None
# Add shortcut if number of layers is larger than 1
if not single_layer:
# Use identity is input and output reprs are identical
if in_repr == out_reprs[-1] and stride == 1:
self.shortcut = lambda x: x
else:
self.shortcut = []
self.shortcut.append(
GatedBlock(reprs[0], reprs[-1],
size=size, padding=size//2,
stride=conv_stride,
activation=None,
radial_window=radial_window,
batch_norm_momentum=batch_norm_momentum,
normalization=normalization,
capsule_dropout_p=capsule_dropout_p))
if downsample_by_pooling and stride > 1:
self.shortcut.append(nn.AvgPool3d(kernel_size=size,
padding=size//2,
stride=stride))
self.shortcut = nn.Sequential(*self.shortcut)
self.activation = GatedActivation(
repr_in=reprs[-1],
size=size,
radial_window=radial_window,
batch_norm_momentum=batch_norm_momentum,
normalization=normalization)
def __init__(self, n_out):
super().__init__()
self.int_repr = None
features = [
(1, ), # 64
(8, 4, 2),
(8, 4, 2), # 34, 38
(16, 8, 4),
(16, 8, 4), # 21, 25
(32, 16, 8),
(32, 16, 8), # 15, 19
(32, 16, 8), # 12
(512, )
]
common_block_params = {
'size': 5,
'padding': 4,
'normalization': 'batch',
'smooth_stride': True,
'capsule_dropout_p': 0.1,
'dyn_iso': True,
'radial_window': se3cnn.kernel.sigmoid_window,
}
block_params = [
{
'activation': (F.relu, torch.sigmoid),
'stride': 2
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid),
'stride': 2
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid),
'stride': 2
},
{
'activation': (F.relu, torch.sigmoid)
},
{
'activation': (F.relu, torch.sigmoid),
'stride': 2
},
{
'activation': (F.relu, torch.sigmoid)
},
]
assert len(block_params) + 1 == len(features)
blocks = [
GatedBlock(features[i], features[i + 1], **common_block_params,
**block_params[i]) for i in range(len(block_params))
]
self.sequence = torch.nn.Sequential(
*blocks,
AvgSpacial(),
nn.Linear(features[-1][0], n_out),
)