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)