def __init__(self, num_classes):
super(Net, self).__init__()
self.conv1 = XConv(0, 48, dim=3, kernel_size=8, hidden_channels=32)
self.conv2 = XConv(48,
96,
dim=3,
kernel_size=12,
hidden_channels=64,
dilation=2)
self.conv3 = XConv(96,
192,
dim=3,
kernel_size=16,
hidden_channels=128,
dilation=2)
self.conv4 = XConv(192,
384,
dim=3,
kernel_size=16,
hidden_channels=256,
dilation=2)
self.lin1 = Lin(384, 256)
self.lin2 = Lin(256, 128)
self.lin3 = Lin(128, num_classes)
def __init__(self, NUM_CLASS):
super(RotationEstimateNN, self).__init__()
self.NUM_CLASS = NUM_CLASS
self.cls_model = Classifier(NUM_CLASS)
self.cls_model.load_state_dict(torch.load(os.path.join(BASE_DIR, 'models', 'classifier_1024', '117_0.98_1024')))
self.cls_model.eval()
self.pcnn1 = XConv(0, 48, dim=3, kernel_size=8, hidden_channels=32)
# in_c , out_c , nei_c , spread, n_rep
self.pcnn2 = XConv(
48, 96, dim=3, kernel_size=12, hidden_channels=64, dilation=2)
self.pcnn3 = XConv(
96, 192, dim=3, kernel_size=16, hidden_channels=128, dilation=2)
# self.pcnn4 = XConv(
# 192, 384, dim=3, kernel_size=16, hidden_channels=256, dilation=2)
# self.pcnn5 = XConv(
# 384, 768, dim=3, kernel_size=16, hidden_channels=512, dilation=2)
self.linctr = Lin(3, 200)
self.lin1 = Lin(400, 512)
self.lin2 = Lin(512, 256)
self.lin3 = Lin(256, 128)
self.lin_psi = Lin(128, 12+1)
self.lin_theta = Lin(128, 12+1)
self.lin_phi = Lin(128, 12+1)
def __init__(self, NUM_CLASS):
super(Classifier, self).__init__()
self.pcnn1 = XConv(0, 48, dim=3, kernel_size=8, hidden_channels=32)
# in_c , out_c , nei_c , spread, n_rep
self.pcnn2 = XConv(48,
96,
dim=3,
kernel_size=12,
hidden_channels=64,
dilation=2)
self.pcnn3 = XConv(96,
192,
dim=3,
kernel_size=16,
hidden_channels=128,
dilation=2)
# self.pcnn4 = XConv(
# 192, 384, dim=3, kernel_size=16, hidden_channels=256, dilation=2)
# self.pcnn5 = XConv(
# 384, 768, dim=3, kernel_size=16, hidden_channels=512, dilation=2)
self.lin1 = Lin(192, 256)
self.lin2 = Lin(256, 128)
self.lin3 = Lin(128, NUM_CLASS)
def test_x_conv():
in_channels, out_channels = (16, 32)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
num_nodes = edge_index.max().item() + 1
x = torch.randn((num_nodes, in_channels))
pos = torch.rand((num_nodes, 3))
conv = XConv(in_channels, out_channels, dim=3, kernel_size=2, dilation=2)
assert conv.__repr__() == 'XConv(16, 32)'
assert conv(x, pos).size() == (num_nodes, out_channels)
def __init__(self, pool='mean'):
super().__init__()
self.conv1 = XConv(0, 48, dim=3, kernel_size=8, hidden_channels=32)
self.conv2 = XConv(
48, 96, dim=3, kernel_size=12, hidden_channels=64, dilation=2)
self.conv3 = XConv(
96, 192, dim=3, kernel_size=16, hidden_channels=128, dilation=2)
self.conv4 = XConv(
192, 384, dim=3, kernel_size=16, hidden_channels=256, dilation=2)
self.lin1 = Lin(384, 256)
self.lin2 = Lin(256, 128)
self.lin3 = Lin(128, 40)
self.pool = pool
def __init__(self, num_classes, normal_channel=True):
super(get_model, self).__init__()
if normal_channel:
dim = 6
else:
dim = 3
self.conv1 = XConv(0, 48, dim=dim, kernel_size=8, hidden_channels=32)
self.conv2 = XConv(
48, 96, dim=dim, kernel_size=12, hidden_channels=64, dilation=2)
self.conv3 = XConv(
96, 192, dim=dim, kernel_size=16, hidden_channels=128, dilation=2)
self.conv4 = XConv(
192, 384, dim=dim, kernel_size=16, hidden_channels=256, dilation=2)
self.lin1 = Lin(384, 256)
self.lin2 = Lin(256, 128)
self.lin3 = Lin(128, num_classes)
def test_x_conv():
x = torch.randn(8, 16)
pos = torch.rand(8, 3)
batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1])
conv = XConv(16, 32, dim=3, kernel_size=2, dilation=2)
assert conv.__repr__() == 'XConv(16, 32)'
torch.manual_seed(12345)
out1 = conv(x, pos)
assert out1.size() == (8, 32)
torch.manual_seed(12345)
out2 = conv(x, pos, batch)
assert out2.size() == (8, 32)
jit = torch.jit.script(conv)
torch.manual_seed(12345)
assert jit(x, pos).tolist() == out1.tolist()
torch.manual_seed(12345)
assert jit(x, pos, batch).tolist() == out2.tolist()
def __init__(self, d=3):
super(PointCNNPolicyNetwork, self).__init__()
self.conv1 = XConv(d, 16, dim=2, kernel_size=10, hidden_channels=4)
self.conv2 = XConv(16, 16, dim=2, kernel_size=10, hidden_channels=4)
self.conv3 = XConv(16, 1, dim=2, kernel_size=10, hidden_channels=4)
self.fc = nn.Linear(16, 1)