def __init__(self, in_channels, out_channels, laplacian):
super().__init__()
self.conv1 = ConvCheb(in_channels, out_channels, 1, laplacian)
self.conv2 = ConvCheb(out_channels, out_channels, 3, laplacian)
self.conv3 = ConvCheb(out_channels, in_channels, 1, laplacian)
self.bn1 = BatchNorm1d(out_channels)
self.bn2 = BatchNorm1d(out_channels)
self.bn3 = BatchNorm1d(in_channels)
def __init__(self, N, in_channels, out_channels, kernel_size):
super().__init__()
self.kernel_size = kernel_size
laplacians = []
for i, nodes in enumerate([3072, 768, 192]):
laplacian = _compute_laplacian_healpix(nodes)
laplacians.append(laplacian)
# Pooling - unpooling
self.pooling = PoolMaxHealpix(kernel_size=4)
self.unpool = UnpoolMaxHealpix(kernel_size=4)
# Encoding block 1
self.conv11 = ConvBlock(in_channels, max(in_channels, 64), kernel_size,
laplacians[0])
self.conv12 = ConvBlock(max(in_channels, 64), max(in_channels, 64),
kernel_size, laplacians[0])
self.conv13 = ConvCheb(
max(in_channels, 64), max(in_channels,
64), kernel_size, laplacians[0]
) #ConvBlock(max(in_channels, 32), 64, kernel_size, laplacians[0])
self.conv1_res = Conv1dAuto(in_channels, 64, 1)
# Encoding block 2
self.conv21 = ConvBlock(64, 128, kernel_size, laplacians[1])
self.conv22 = ConvBlock(128, 128, kernel_size, laplacians[1])
self.conv23 = ConvCheb(128, 128, kernel_size, laplacians[1])
self.conv2_res = Conv1dAuto(64, 128, 1)
# Encoding block 3
self.conv31 = ConvBlock(128, 256, kernel_size, laplacians[2])
self.conv32 = ConvBlock(256, 256, kernel_size, laplacians[2])
self.conv33 = ConvCheb(256, 128, kernel_size, laplacians[2])
self.conv3_res = Conv1dAuto(128, 128, 1)
# Decoding block 4
self.uconv21 = ConvBlock(256, 128, kernel_size, laplacians[1])
self.uconv22 = ConvBlock(128, 64, kernel_size, laplacians[1])
# Decoding block 4
self.uconv11 = ConvBlock(128, 64, kernel_size, laplacians[0])
self.uconv12 = ConvBlock(64, 32, kernel_size, laplacians[0])
self.uconv13 = ConvCheb(32, out_channels, kernel_size, laplacians[0])
def __init__(self, N, ratio, in_channels, out_channels, kernel_size):
super().__init__()
self.kernel_size = kernel_size
laplacian = _compute_laplacian_equiangular(N, ratio)
self.register_buffer(f'laplacian', laplacian)
self.conv1 = ConvCheb(in_channels, 64, self.kernel_size)
self.conv2 = ConvCheb(64, 64, self.kernel_size)
self.conv3 = ConvCheb(64, 64, self.kernel_size)
self.conv4 = ConvCheb(64, 64, self.kernel_size)
self.conv5 = ConvCheb(64, out_channels, self.kernel_size)
for m in self.modules():
if isinstance(m, ConvCheb):
m.reset_parameters(activation='linear',
fan='avg',
distribution='uniform')
def __init__(self,
N,
in_channels,
out_channels,
kernel_size,
kernel_size_pooling=4):
super().__init__()
self.kernel_size = kernel_size
num_nodes = 3072
laplacians = []
for i, nodes in enumerate([
num_nodes, num_nodes / kernel_size_pooling,
num_nodes / (kernel_size * kernel_size_pooling)
]):
laplacian = _compute_laplacian_healpix(nodes)
laplacians.append(laplacian)
# Pooling - unpooling
self.pooling = PoolMaxHealpix(kernel_size=kernel_size_pooling)
self.unpool = UnpoolMaxHealpix(kernel_size=kernel_size_pooling)
# Encoding block 1
self.conv11 = ConvBlock(in_channels, max(in_channels, 32 * 2),
kernel_size, laplacians[0])
self.conv13 = ConvBlock(max(in_channels, 32 * 2), 64 * 2, kernel_size,
laplacians[0])
self.conv1_res = Conv1dAuto(in_channels, 64 * 2, 1)
# Encoding block 2
self.conv21 = ConvBlock(64 * 2, 96 * 2, kernel_size, laplacians[1])
self.conv23 = ConvBlock(96 * 2, 128 * 2, kernel_size, laplacians[1])
self.conv2_res = Conv1dAuto(64 * 2, 128 * 2, 1)
# Encoding block 3
self.conv31 = ConvBlock(128 * 2, 256 * 2, kernel_size, laplacians[2])
# self.conv32 = ConvBlock(256*2, 256*2, kernel_size, laplacians[2])
self.conv33 = ConvBlock(256 * 2, 128 * 2, kernel_size, laplacians[2])
self.conv3_res = Conv1dAuto(128 * 2, 128 * 2, 1)
# Decoding block 4
self.uconv21 = ConvBlock(256 * 2, 128 * 2, kernel_size, laplacians[1])
self.uconv22 = ConvBlock(128 * 2, 64 * 2, kernel_size, laplacians[1])
# Decoding block 4
self.uconv11 = ConvBlock(128 * 2, 64 * 2, kernel_size, laplacians[0])
self.uconv12 = ConvBlock(64 * 2, 32 * 2, kernel_size, laplacians[0])
self.uconv13 = ConvCheb(32 * 2, out_channels, kernel_size,
laplacians[0])
def __init__(self, N, ratio, in_channels, out_channels, kernel_size):
super().__init__()
self.ratio = ratio
self.kernel_size = kernel_size
for i, nodes in enumerate([8, 32, 128, 512, 2048]):
laplacian = _compute_laplacian(nodes, ratio)
self.register_buffer(f'laplacian_{i+1}', laplacian)
# Pooling - unpooling
self.pooling = PoolAvgEquiangular(ratio, 2)
self.unpool = UnpoolAvgEquiangular(ratio, 2)
# Encoding block 1
self.conv1_enc_l5 = ConvCheb(in_channels, 32, self.kernel_size)
self.conv2_enc_l5 = ConvCheb(32, 64, self.kernel_size)
# Encoding block 2
self.conv_enc_l4 = ConvCheb(64, 128, self.kernel_size)
# Encoding block 3
self.conv_enc_l3 = ConvCheb(128, 256, self.kernel_size)
# Encoding block 4
self.conv_enc_l2 = ConvCheb(256, 512, self.kernel_size)
# Encoding block 5
self.conv_enc_l1 = ConvCheb(512, 512, self.kernel_size)
# Encoding block 6
self.conv_enc_l0 = ConvCheb(512, 512, self.kernel_size)
# Decoding block 1
self.conv1_dec_l1 = ConvCheb(512, 512, self.kernel_size)
self.conv2_dec_l1 = ConvCheb(512 + 512, 512, self.kernel_size)
# Decoding block 2
self.conv1_dec_l2 = ConvCheb(512, 256, self.kernel_size)
self.conv2_dec_l2 = ConvCheb(512 + 256, 256, self.kernel_size)
# Decoding block 3
self.conv1_dec_l3 = ConvCheb(256, 128, self.kernel_size)
self.conv2_dec_l3 = ConvCheb(256 + 128, 128, self.kernel_size)
# Decoding block 4
self.conv1_dec_l4 = ConvCheb(128, 64, self.kernel_size)
self.conv2_dec_l4 = ConvCheb(128 + 64, 64, self.kernel_size)
# Decoding block 5
self.conv1_dec_l5 = ConvCheb(64, 32, self.kernel_size)
self.conv2_dec_l5 = ConvCheb(32, out_channels, self.kernel_size)