def __init__(self, output_channels=40):
super(Point_Transformer, self).__init__()
self.conv1 = nn.Conv1d(3, 128, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(128, 128, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(128)
self.sa1 = SA_Layer(128)
self.sa2 = SA_Layer(128)
self.sa3 = SA_Layer(128)
self.sa4 = SA_Layer(128)
self.conv_fuse = nn.Sequential(
nn.Conv1d(512, 1024, kernel_size=1, bias=False),
nn.BatchNorm1d(1024), nn.LeakyReLU(scale=0.2))
self.linear1 = nn.Linear(1024, 512, bias=False)
self.bn6 = nn.BatchNorm1d(512)
self.dp1 = nn.Dropout(p=0.5)
self.linear2 = nn.Linear(512, 256)
self.bn7 = nn.BatchNorm1d(256)
self.dp2 = nn.Dropout(p=0.5)
self.linear3 = nn.Linear(256, output_channels)
self.relu = nn.ReLU()
def __init__(self, in_channels, out_channels):
super(Local_op, self).__init__()
self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(out_channels)
self.bn2 = nn.BatchNorm1d(out_channels)
self.relu = nn.ReLU()
def __init__(self, n_classes=40):
super(DGCNN, self).__init__()
self.k = 20
self.knn = KNN(self.k)
self.bn1 = nn.BatchNorm(64)
self.bn2 = nn.BatchNorm(64)
self.bn3 = nn.BatchNorm(128)
self.bn4 = nn.BatchNorm(256)
self.bn5 = nn.BatchNorm1d(1024)
self.conv1 = nn.Sequential(nn.Conv(6, 64, kernel_size=1, bias=False),
self.bn1,
nn.LeakyReLU(scale=0.2))
self.conv2 = nn.Sequential(nn.Conv(64*2, 64, kernel_size=1, bias=False),
self.bn2,
nn.LeakyReLU(scale=0.2))
self.conv3 = nn.Sequential(nn.Conv(64*2, 128, kernel_size=1, bias=False),
self.bn3,
nn.LeakyReLU(scale=0.2))
self.conv4 = nn.Sequential(nn.Conv(128*2, 256, kernel_size=1, bias=False),
self.bn4,
nn.LeakyReLU(scale=0.2))
self.conv5 = nn.Sequential(nn.Conv1d(512, 1024, kernel_size=1, bias=False),
self.bn5,
nn.LeakyReLU(scale=0.2))
self.linear1 = nn.Linear(1024*2, 512, bias=False)
self.bn6 = nn.BatchNorm1d(512)
self.dp1 = nn.Dropout(p=0.5)
self.linear2 = nn.Linear(512, 256)
self.bn7 = nn.BatchNorm1d(256)
self.dp2 = nn.Dropout(p=0.5)
self.linear3 = nn.Linear(256, n_classes)
def __init__(self, n_classes=40):
super(PointConvDensityClsSsg, self).__init__()
self.sa1 = PointConvDensitySetAbstraction(npoint=512,
nsample=32,
in_channel=3,
mlp=[64, 64, 128],
bandwidth=0.1,
group_all=False)
self.sa2 = PointConvDensitySetAbstraction(npoint=128,
nsample=64,
in_channel=128 + 3,
mlp=[128, 128, 256],
bandwidth=0.2,
group_all=False)
self.sa3 = PointConvDensitySetAbstraction(npoint=1,
nsample=None,
in_channel=256 + 3,
mlp=[256, 512, 1024],
bandwidth=0.4,
group_all=True)
self.fc1 = nn.Linear(1024, 512)
self.bn1 = nn.BatchNorm1d(512)
self.drop1 = nn.Dropout(0.4)
self.fc2 = nn.Linear(512, 256)
self.bn2 = nn.BatchNorm1d(256)
self.drop2 = nn.Dropout(0.4)
self.fc3 = nn.Linear(256, n_classes)
self.relu = nn.ReLU()
def __init__(self, cin, cout, num_points, num_neighbors):
super(KNNEmbed, self).__init__()
self.num_points = num_points
self.num_neighbors = num_neighbors
self.embeds = nn.Sequential(
nn.Conv1d(cin, cout, kernel_size=1, bias=False),
nn.BatchNorm1d(cout), nn.ReLU(),
nn.Conv1d(cout, cout, kernel_size=1, bias=False),
nn.BatchNorm1d(cout), nn.ReLU())
def __init__(self, channels=256):
super(Point_Transformer_Last, self).__init__()
self.conv1 = nn.Conv1d(channels, channels, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(channels, channels, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(channels)
self.bn2 = nn.BatchNorm1d(channels)
self.sa1 = SA_Layer(channels)
self.sa2 = SA_Layer(channels)
self.sa3 = SA_Layer(channels)
self.sa4 = SA_Layer(channels)
self.relu = nn.ReLU()
def __init__(self):
super(Discriminator, self).__init__()
self.down = nn.Sequential(nn.Conv(opt.channels, 64, 3, stride=2, padding=1), nn.ReLU())
self.down_size = (opt.img_size // 2)
down_dim = (64 * ((opt.img_size // 2) ** 2))
self.embedding = nn.Linear(down_dim, 32)
self.fc = nn.Sequential(
nn.BatchNorm1d(32, 0.8),
nn.ReLU(),
nn.Linear(32, down_dim),
nn.BatchNorm1d(down_dim),
nn.ReLU()
)
self.up = nn.Sequential(nn.Upsample(scale_factor=2), nn.Conv(64, opt.channels, 3, stride=1, padding=1))
def __init__(self, part_num):
super(DGCNN_partseg, self).__init__()
self.seg_num_all = part_num
self.k = 40
self.knn = KNN(self.k)
self.bn1 = nn.BatchNorm2d(64)
self.bn2 = nn.BatchNorm2d(64)
self.bn3 = nn.BatchNorm2d(64)
self.bn4 = nn.BatchNorm2d(64)
self.bn5 = nn.BatchNorm2d(64)
self.bn6 = nn.BatchNorm1d(1024)
self.bn7 = nn.BatchNorm1d(64)
self.bn8 = nn.BatchNorm1d(256)
self.bn9 = nn.BatchNorm1d(256)
self.bn10 = nn.BatchNorm1d(128)
self.conv1 = nn.Sequential(nn.Conv2d(6, 64, kernel_size=1, bias=False),
self.bn1,
nn.LeakyReLU(scale=0.2))
self.conv2 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False),
self.bn2,
nn.LeakyReLU(scale=0.2))
self.conv3 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
self.bn3,
nn.LeakyReLU(scale=0.2))
self.conv4 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=1, bias=False),
self.bn4,
nn.LeakyReLU(scale=0.2))
self.conv5 = nn.Sequential(nn.Conv2d(64*2, 64, kernel_size=1, bias=False),
self.bn5,
nn.LeakyReLU(scale=0.2))
self.conv6 = nn.Sequential(nn.Conv1d(192, 1024, kernel_size=1, bias=False),
self.bn6,
nn.LeakyReLU(scale=0.2))
self.conv7 = nn.Sequential(nn.Conv1d(16, 64, kernel_size=1, bias=False),
self.bn7,
nn.LeakyReLU(scale=0.2))
self.conv8 = nn.Sequential(nn.Conv1d(1280, 256, kernel_size=1, bias=False),
self.bn8,
nn.LeakyReLU(scale=0.2))
self.dp1 = nn.Dropout(p=0.5)
self.conv9 = nn.Sequential(nn.Conv1d(256, 256, kernel_size=1, bias=False),
self.bn9,
nn.LeakyReLU(scale=0.2))
self.dp2 = nn.Dropout(p=0.5)
self.conv10 = nn.Sequential(nn.Conv1d(256, 128, kernel_size=1, bias=False),
self.bn10,
nn.LeakyReLU(scale=0.2))
self.conv11 = nn.Conv1d(128, self.seg_num_all, kernel_size=1, bias=False)
def __init__(self, hidden_unit=[8, 8]):
super(DensityNet, self).__init__()
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
self.mlp_convs.append(nn.Conv1d(1, hidden_unit[0], 1))
self.mlp_bns.append(nn.BatchNorm1d(hidden_unit[0]))
for i in range(1, len(hidden_unit)):
self.mlp_convs.append(
nn.Conv1d(hidden_unit[i - 1], hidden_unit[i], 1))
self.mlp_bns.append(nn.BatchNorm1d(hidden_unit[i]))
self.mlp_convs.append(nn.Conv1d(hidden_unit[-1], 1, 1))
self.mlp_bns.append(nn.BatchNorm1d(1))
self.sigmoid = nn.Sigmoid()
self.relu = nn.ReLU()
def __init__(self):
super(STN3d, self).__init__()
self.conv1 = nn.Conv1d(3, 64, 1)
self.conv2 = nn.Conv1d(64, 128, 1)
self.conv3 = nn.Conv1d(128, 1024, 1)
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 9)
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
def __init__(self, args, margs):
super(PointCloudTransformer, self).__init__(args, margs)
self.input_embeds = nn.Sequential(
Permute(0, 2, 1), nn.Conv1d(3, 64, kernel_size=1, bias=False),
nn.BatchNorm1d(64), nn.ReLU(),
nn.Conv1d(64, 64, kernel_size=1, bias=False), nn.BatchNorm1d(64),
nn.ReLU(), Permute(0, 2, 1))
self.knn_embeds = nn.Sequential(KNNEmbed(128, 128, 512, 32),
KNNEmbed(256, 256, 256, 32))
self.transformer = PointTransformer()
self.classifier = nn.Sequential(nn.Linear(1024, 512),
nn.BatchNorm1d(512), nn.ReLU(),
nn.Dropout(p=0.5), nn.Linear(512, 256),
nn.BatchNorm1d(256), nn.Dropout(p=0.5),
nn.Linear(256, 40))
def test_batchnorm(self):
# ***************************************************************
# Test BatchNorm Layer
# ***************************************************************
arr = np.random.randn(16,10,224,224)
check_equal_with_istrain(arr, jnn.BatchNorm(10, is_train=True), tnn.BatchNorm2d(10))
class Model(tnn.Module):
def __init__(self):
super(Model, self).__init__()
self.layer = tnn.BatchNorm2d(10)
def forward(self, x):
return self.layer(x)
model = Model()
model.eval()
check_equal_with_istrain(arr, jnn.BatchNorm(10, is_train=False), model, False)
# ***************************************************************
# Test InstanceNorm2d Layer
# ***************************************************************
arr = np.random.randn(16,10,224,224)
check_equal_without_istrain(arr, jnn.InstanceNorm2d(10, is_train=True), tnn.InstanceNorm2d(10))
class Model(tnn.Module):
def __init__(self):
super(Model, self).__init__()
self.layer = tnn.InstanceNorm2d(10)
def forward(self, x):
return self.layer(x)
model = Model()
model.eval()
check_equal_without_istrain(arr, jnn.InstanceNorm2d(10, is_train=False), model)
# ***************************************************************
# Test BatchNorm1d Layer
# ***************************************************************
arr = np.random.randn(16,10)
check_equal_with_istrain(arr, jnn.BatchNorm1d(10, is_train=True), tnn.BatchNorm1d(10), 1e-3)
class Model(tnn.Module):
def __init__(self):
super(Model, self).__init__()
self.layer = tnn.BatchNorm1d(10)
def forward(self, x):
return self.layer(x)
model = Model()
model.eval()
check_equal_with_istrain(arr, jnn.BatchNorm1d(10, is_train=False), model, False)
def __init__(self):
super(Encoder, self).__init__()
self.model = nn.Sequential(nn.Linear(int(np.prod(img_shape)), 512),
nn.Leaky_relu(0.2), nn.Linear(512, 512),
nn.BatchNorm1d(512), nn.Leaky_relu(0.2))
self.mu = nn.Linear(512, opt.latent_dim)
self.logvar = nn.Linear(512, opt.latent_dim)
def __init__(self, part_num=50):
super(PointConvDensity_partseg, self).__init__()
self.part_num = part_num
self.sa0 = PointConvDensitySetAbstraction(npoint=1024, nsample=32, in_channel=3, mlp=[32,32,64], bandwidth = 0.1, group_all=False)
self.sa1 = PointConvDensitySetAbstraction(npoint=256, nsample=32, in_channel=64 + 3, mlp=[64,64,128], bandwidth = 0.2, group_all=False)
self.sa2 = PointConvDensitySetAbstraction(npoint=64, nsample=32, in_channel=128 + 3, mlp=[128,128,256], bandwidth = 0.4, group_all=False)
self.sa3 = PointConvDensitySetAbstraction(npoint=36, nsample=32, in_channel=256 + 3, mlp=[256,256,512], bandwidth = 0.8, group_all=False)
# TODO upsample
# upsampling
# def __init__(self, nsample, in_channel, mlp, bandwidth):
self.in0 = PointConvDensitySetInterpolation(nsample=16, in_channel=512 + 3, mlp=[512,512], bandwidth=0.8)
self.in1 = PointConvDensitySetInterpolation(nsample=16, in_channel=512 + 3, mlp=[256,256], bandwidth=0.4)
self.in2 = PointConvDensitySetInterpolation(nsample=16, in_channel=256 + 3, mlp=[128,128], bandwidth=0.2)
self.in3 = PointConvDensitySetInterpolation(nsample=16, in_channel=128 + 3, mlp=[128,128, 128], bandwidth=0.1)
# self.fp0 = PointConvDensitySetAbstraction(npoint=1024, nsample=32, in_channel=3, mlp=[32,32,64], bandwidth = 0.1, group_all=False)
# self.fp1 = PointConvDensitySetAbstraction(npoint=256, nsample=32, in_channel=64 + 3, mlp=[64,64,128], bandwidth = 0.2, group_all=False)
# self.fp2 = PointConvDensitySetAbstraction(npoint=64, nsample=32, in_channel=128 + 3, mlp=[128,128,256], bandwidth = 0.4, group_all=False)
# self.fp3 = PointConvDensitySetAbstraction(npoint=36, nsample=32, in_channel=256 + 3, mlp=[256,256,512], bandwidth = 0.8, group_all=False)
self.fc1 = nn.Conv1d(128, 128, 1)
self.bn1 = nn.BatchNorm1d(128)
self.drop1 = nn.Dropout(0.4)
self.fc3 = nn.Conv1d(128, self.part_num, 1)
self.relu = nn.ReLU()
def __init__(self):
super(Decoder, self).__init__()
self.model = nn.Sequential(nn.Linear(opt.latent_dim, 512),
nn.Leaky_relu(0.2), nn.Linear(512, 512),
nn.BatchNorm1d(512), nn.Leaky_relu(0.2),
nn.Linear(512, int(np.prod(img_shape))),
nn.Tanh())
def build_model(self):
self.pointnet_modules = nn.ModuleList()
self.pointnet_modules.append(
PointnetModule(
n_points=512,
radius=0.2,
n_samples=64,
mlp=[3, 64, 64, 128],
use_xyz=self.use_xyz,
))
self.pointnet_modules.append(
PointnetModule(
n_points=128,
radius=0.4,
n_samples=64,
mlp=[128, 128, 128, 256],
use_xyz=self.use_xyz,
))
self.pointnet_modules.append(
PointnetModule(
mlp=[256, 256, 512, 1024],
use_xyz=self.use_xyz,
))
self.fp3 = PointNetFeaturePropagation(in_channel=1280, mlp=[256, 256])
self.fp2 = PointNetFeaturePropagation(in_channel=384, mlp=[256, 128])
self.fp1 = PointNetFeaturePropagation(in_channel=128 + 16 + 6,
mlp=[128, 128, 128])
self.fc_layer = nn.Sequential(nn.Conv1d(128, 128, 1),
nn.BatchNorm1d(128), nn.Dropout(0.5),
nn.Conv1d(128, self.part_num, 1))
def __init__(self, in_channel, out_channel, kernel_size=1, stride=1, if_bn=True, activation_fn=nn.Relu()):
super(Conv1d, self).__init__()
self.conv = nn.Conv1d(in_channel, out_channel, kernel_size, stride=stride)
# self.conv = nn.Linear(in_channel, out_channel)
self.if_bn = if_bn
if self.if_bn:
self.bn = nn.BatchNorm1d(out_channel)
self.activation_fn = activation_fn
def __init__(self, in_channel, mlp):
super(PointNetFeaturePropagation, self).__init__()
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
self.relu = nn.ReLU()
for out_channel in mlp:
self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm1d(out_channel))
last_channel = out_channel
def __init__(self, channels):
super(SA_Layer, self).__init__()
self.q_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
self.k_conv = nn.Conv1d(channels, channels // 4, 1, bias=False)
self.q_conv.conv.weight = self.k_conv.conv.weight
self.v_conv = nn.Conv1d(channels, channels, 1)
self.trans_conv = nn.Conv1d(channels, channels, 1)
self.after_norm = nn.BatchNorm1d(channels)
self.act = nn.ReLU()
self.softmax = nn.Softmax(dim=-1)
def __init__(self, channels=256):
super(PointTransformer, self).__init__()
self.conv = nn.Sequential(
nn.Conv1d(channels, channels, kernel_size=1, bias=False),
nn.BatchNorm1d(channels), nn.ReLU(),
nn.Conv1d(channels, channels, kernel_size=1, bias=False),
nn.BatchNorm1d(channels), nn.ReLU())
self.pos_conv = nn.Sequential(
Permute(0, 2, 1), nn.Conv1d(3, channels, kernel_size=1,
bias=False))
self.oa1 = OffsetAttention(channels)
self.oa2 = OffsetAttention(channels)
self.oa3 = OffsetAttention(channels)
self.oa4 = OffsetAttention(channels)
self.fuse_conv = nn.Sequential(
nn.Conv1d(1280, 1024, kernel_size=1, bias=False),
nn.BatchNorm1d(1024), nn.LeakyReLU(scale=0.2, ))
def build_model(self):
self.pointnet_modules = nn.ModuleList()
self.pointnet_modules.append(
PointnetModule(
n_points=512,
radius=0.2,
n_samples=64,
mlp=[3, 64, 64, 128],
use_xyz=self.use_xyz,
)
)
self.pointnet_modules.append(
PointnetModule(
n_points=128,
radius=0.4,
n_samples=64,
mlp=[128, 128, 128, 256],
use_xyz=self.use_xyz,
)
)
self.pointnet_modules.append(
PointnetModule(
mlp=[256, 256, 512, 1024],
use_xyz=self.use_xyz,
)
)
self.fc_layer = nn.Sequential(
nn.Linear(1024, 512, bias=False),
nn.BatchNorm1d(512),
nn.ReLU(),
nn.Linear(512, 256, bias=False),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, self.n_classes),
)
def __init__(self, part_num=50):
super(Point_Transformer_partseg, self).__init__()
self.part_num = part_num
self.conv1 = nn.Conv1d(3, 128, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(128, 128, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(128)
self.sa1 = SA_Layer(128)
self.sa2 = SA_Layer(128)
self.sa3 = SA_Layer(128)
self.sa4 = SA_Layer(128)
self.conv_fuse = nn.Sequential(
nn.Conv1d(512, 1024, kernel_size=1, bias=False),
nn.BatchNorm1d(1024), nn.LeakyReLU(scale=0.2))
self.label_conv = nn.Sequential(
nn.Conv1d(16, 64, kernel_size=1, bias=False), nn.BatchNorm1d(64),
nn.LeakyReLU(scale=0.2))
self.convs1 = nn.Conv1d(1024 * 3 + 64, 512, 1)
self.dp1 = nn.Dropout(0.5)
self.convs2 = nn.Conv1d(512, 256, 1)
self.convs3 = nn.Conv1d(256, self.part_num, 1)
self.bns1 = nn.BatchNorm1d(512)
self.bns2 = nn.BatchNorm1d(256)
self.relu = nn.ReLU()
def __init__(self, output_channels=40):
super(Point_Transformer2, self).__init__()
self.conv1 = nn.Conv1d(3, 64, kernel_size=1, bias=False)
self.conv2 = nn.Conv1d(64, 64, kernel_size=1, bias=False)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(64)
self.gather_local_0 = Local_op(in_channels=128, out_channels=128)
self.gather_local_1 = Local_op(in_channels=256, out_channels=256)
self.pt_last = Point_Transformer_Last()
self.relu = nn.ReLU()
self.conv_fuse = nn.Sequential(
nn.Conv1d(1280, 1024, kernel_size=1, bias=False),
nn.BatchNorm1d(1024), nn.LeakyReLU(scale=0.2))
self.linear1 = nn.Linear(1024, 512, bias=False)
self.bn6 = nn.BatchNorm1d(512)
self.dp1 = nn.Dropout(p=0.5)
self.linear2 = nn.Linear(512, 256)
self.bn7 = nn.BatchNorm1d(256)
self.dp2 = nn.Dropout(p=0.5)
self.linear3 = nn.Linear(256, output_channels)
def __init__(self):
super(Discriminator, self).__init__()
self.down = nn.Sequential(
nn.Conv(opt.channels, 64, 3, 2, 1),
nn.Relu()
)
self.down_size = (opt.img_size // 2)
down_dim = (64 * ((opt.img_size // 2) ** 2))
self.fc = nn.Sequential(
nn.Linear(down_dim, 32),
nn.BatchNorm1d(32, 0.8),
nn.Relu(),
nn.Linear(32, down_dim),
nn.BatchNorm1d(down_dim),
nn.Relu()
)
self.up = nn.Sequential(
nn.Upsample(scale_factor=2),
nn.Conv(64, opt.channels, 3, 1, 1)
)
for m in self.modules():
weights_init_normal(m)
def __init__(self, in_features : int, out_features : int,
drop_rate : int = 0, with_bn : bool = True,
act = nn.ReLU()
) -> None:
"""
:param in_features: Length of input featuers (last dimension).
:param out_features: Length of output features (last dimension).
:param drop_rate: Drop rate to be applied after activation.
:param with_bn: Whether or not to apply batch normalization.
:param activation: Activation function.
"""
super(DenseConv1d, self).__init__()
self.linear = nn.Conv1d(in_features, out_features, 1)
self.act = act
self.with_bn = with_bn
self.drop = nn.Dropout(drop_rate) if drop_rate > 0 else None
self.bn = nn.BatchNorm1d(out_features) if with_bn else None
def __init__(self, nsample, in_channel, mlp, bandwidth):
super(PointConvDensitySetInterpolation, self).__init__()
self.bandwidth = bandwidth
self.nsample = nsample
self.in_channel = in_channel
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
self.relu = nn.ReLU()
last_channel = in_channel
self.weightnet = WeightNet(3, 16)
self.densitynet = DensityNet()
for out_channel in mlp:
self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm2d(out_channel))
last_channel = out_channel
self.linear = nn.Linear(16 * mlp[-1], mlp[-1])
self.bn_linear = nn.BatchNorm1d(mlp[-1])
def __init__(self, in_channel, mlp):
super(PointNetFeaturePropagation, self).__init__()
# self.mlp_convs = nn.ModuleList()
# self.mlp_bns = nn.ModuleList()
last_channel = in_channel
# self.relu = nn.ReLU()
arr_mlp = []
for out_channel in mlp:
arr_mlp.append(nn.Conv1d(last_channel, out_channel, 1))
arr_mlp.append(nn.BatchNorm1d(out_channel))
arr_mlp.append(nn.ReLU())
last_channel = out_channel
# for out_channel in mlp:
# self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
# self.mlp_bns.append(nn.BatchNorm1d(out_channel))
# last_channel = out_channel
self.mlp = nn.Sequential(*arr_mlp)
def __init__(self, channels):
super(OffsetAttention, self).__init__()
self.q_conv = nn.Conv1d(channels,
channels // 4,
kernel_size=1,
bias=False)
self.k_conv = nn.Conv1d(channels,
channels // 4,
kernel_size=1,
bias=False)
self.q_conv.weight = self.k_conv.weight
self.v_conv = nn.Conv1d(channels, channels, kernel_size=1, bias=False)
self.trans_conv = nn.Conv1d(channels,
channels,
kernel_size=1,
bias=False)
self.after_norm = nn.BatchNorm1d(channels)
self.act = nn.ReLU()
self.softmax = nn.Softmax(dim=-1)
def __init__(self, npoint, nsample, in_channel, mlp, bandwidth, group_all):
super(PointConvDensitySetAbstraction, self).__init__()
self.npoint = npoint
self.nsample = nsample
self.mlp_convs = nn.ModuleList()
self.mlp_bns = nn.ModuleList()
last_channel = in_channel
for out_channel in mlp:
self.mlp_convs.append(nn.Conv(last_channel, out_channel, 1))
self.mlp_bns.append(nn.BatchNorm(out_channel))
last_channel = out_channel
self.weightnet = WeightNet(3, 16)
self.densitynet = DensityNet()
self.linear = nn.Linear(16 * mlp[-1], mlp[-1])
self.bn_linear = nn.BatchNorm1d(mlp[-1])
self.group_all = group_all
self.bandwidth = bandwidth
self.relu = nn.ReLU()
def __init__(self, part_num=50):
super(PointNet_partseg, self).__init__()
self.part_num = part_num
self.stn = STN3d()
self.conv1 = nn.Conv1d(3, 64, 1)
self.conv2 = nn.Conv1d(64, 128, 1)
self.conv3 = nn.Conv1d(128, 128, 1)
self.conv4 = nn.Conv1d(128, 512, 1)
self.conv5 = nn.Conv1d(512, 2048, 1)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(128)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(2048)
self.fstn = STNkd(k=128)
self.convs1 = nn.Conv1d(4944, 256, 1)
self.convs2 = nn.Conv1d(256, 256, 1)
self.convs3 = nn.Conv1d(256, 128, 1)
self.convs4 = nn.Conv1d(128, part_num, 1)
self.bns1 = nn.BatchNorm1d(256)
self.bns2 = nn.BatchNorm1d(256)
self.bns3 = nn.BatchNorm1d(128)
self.relu = nn.ReLU()
