learning3d 0.0.1__py3-none-any.whl

learning3d/losses/cuda/emd_torch/pkg/include/emd.h

@@ -0,0 +1,54 @@
+#ifndef EMD_H_
+#define EMD_H_
+
+#include <torch/extension.h>
+#include <vector>
+
+#include "cuda_helper.h"
+
+
+std::vector<at::Tensor> emd_forward_cuda(
+	at::Tensor xyz1,
+	at::Tensor xyz2);
+
+std::vector<at::Tensor> emd_backward_cuda(
+	at::Tensor xyz1, 
+	at::Tensor xyz2, 
+	at::Tensor match);
+
+// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
+// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
+// CALL FUNCTION IMPLEMENTATIONS
+// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
+// * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
+
+std::vector<at::Tensor> emd_forward(
+	at::Tensor xyz1, 
+	at::Tensor xyz2)
+{
+	CHECK_INPUT(xyz1);
+	CHECK_INPUT(xyz2);
+
+	return emd_forward_cuda(xyz1, xyz2);
+}
+
+std::vector<at::Tensor> emd_backward(
+	at::Tensor xyz1, 
+	at::Tensor xyz2,
+	at::Tensor match)
+{
+	CHECK_INPUT(xyz1);
+	CHECK_INPUT(xyz2);
+	CHECK_INPUT(match);
+
+	return emd_backward_cuda(xyz1, xyz2, match);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+	m.def("emd_forward", &emd_forward, "Compute Earth Mover's Distance");
+	m.def("emd_backward", &emd_backward, "Compute Gradients for Earth Mover's Distance");
+}
+
+
+
+#endif

learning3d/losses/cuda/emd_torch/pkg/layer/__init__.py

@@ -0,0 +1 @@
+from .emd_loss_layer import EMDLoss

learning3d/losses/cuda/emd_torch/pkg/layer/emd_loss_layer.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+
+import _emd_ext._emd as emd
+
+
+class EMDFunction(torch.autograd.Function):
+	@staticmethod
+	def forward(self, xyz1, xyz2):
+		cost, match = emd.emd_forward(xyz1, xyz2)
+		self.save_for_backward(xyz1, xyz2, match)
+		return cost
+
+
+	@staticmethod
+	def backward(self, grad_output):
+		xyz1, xyz2, match = self.saved_tensors
+		grad_xyz1, grad_xyz2 = emd.emd_backward(xyz1, xyz2, match)
+		return grad_xyz1, grad_xyz2
+
+
+
+
+class EMDLoss(nn.Module):
+	'''
+	Computes the (approximate) Earth Mover's Distance between two point sets. 
+
+	IMPLEMENTATION LIMITATIONS:
+	- Double tensors must have <=11 dimensions
+	- Float tensors must have <=23 dimensions
+	This is due to the use of CUDA shared memory in the computation. This shared memory is limited by the hardware to 48kB.
+	'''
+
+	def __init__(self):
+		super(EMDLoss, self).__init__()
+
+	def forward(self, xyz1, xyz2):
+
+		assert xyz1.shape[-1] == xyz2.shape[-1], 'Both point sets must have the same dimensionality'
+		return EMDFunction.apply(xyz1, xyz2)

learning3d/losses/cuda/emd_torch/pkg/src/cuda/emd.cu

@@ -0,0 +1,70 @@
+#include <ATen/ATen.h>
+
+#include <vector>
+
+#include "cuda/emd.cuh"
+
+
+std::vector<at::Tensor> emd_forward_cuda(
+	at::Tensor xyz1, // B x N1 x D
+	at::Tensor xyz2) // B x N2 x D
+{
+	// Some useful values
+	const int batch_size = xyz1.size(0);
+	const int num_pts_1 = xyz1.size(1);
+	const int num_pts_2 = xyz2.size(1);
+
+	// Allocate necessary data structures
+	at::Tensor match = at::zeros({batch_size, num_pts_1, num_pts_2}, 
+		xyz1.options());
+	at::Tensor cost = at::zeros({batch_size}, xyz1.options());
+	at::Tensor temp = at::zeros({batch_size, 2 * (num_pts_1 + num_pts_2)}, 
+		xyz1.options());
+
+	// Find the approximate matching
+	approxmatchLauncher(
+		batch_size, num_pts_1, num_pts_2,
+		xyz1,
+		xyz2, 
+		match,
+		temp
+	);
+
+	// Compute the matching cost
+	matchcostLauncher(
+		batch_size, num_pts_1, num_pts_2, 
+		xyz1,
+		xyz2, 
+		match,
+		cost
+	);
+
+	return {cost, match};
+}
+
+std::vector<at::Tensor> emd_backward_cuda(
+	at::Tensor xyz1, 
+	at::Tensor xyz2, 
+	at::Tensor match)
+{
+	// Some useful values
+	const int batch_size = xyz1.size(0);
+	const int num_pts_1 = xyz1.size(1);
+	const int num_pts_2 = xyz2.size(1);
+
+	// Allocate necessary data structures
+	at::Tensor grad_xyz1 = at::zeros_like(xyz1);
+	at::Tensor grad_xyz2 = at::zeros_like(xyz2);
+
+	// Compute the gradient with respect to the two inputs (xyz1 and xyz2)
+	matchcostgradLauncher(
+		batch_size, num_pts_1, num_pts_2,
+		xyz1,
+		xyz2,
+		match,
+		grad_xyz1,
+		grad_xyz2
+	);
+
+	return {grad_xyz1, grad_xyz2};
+}

learning3d/losses/cuda/emd_torch/pkg/src/emd.cpp

@@ -0,0 +1 @@
+#include "emd.h"

learning3d/losses/cuda/emd_torch/setup.py

@@ -0,0 +1,29 @@
+from setuptools import setup
+from torch.utils.cpp_extension import BuildExtension, CUDAExtension
+
+
+setup(
+	name='PyTorch EMD',
+	version='0.0',
+	author='Vinit Sarode',
+	author_email='vinitsarode5@gmail.com', 
+	description='A PyTorch module for the earth mover\'s distance loss', 
+	ext_package='_emd_ext',
+	ext_modules=[
+		CUDAExtension(
+			name='_emd', 
+			sources=[
+				'pkg/src/emd.cpp',
+				'pkg/src/cuda/emd.cu',
+			],
+			include_dirs=['pkg/include'],
+		),
+	],
+	packages=[
+		'emd', 
+	],
+	package_dir={
+		'emd' : 'pkg/layer'
+	},
+	cmdclass={'build_ext': BuildExtension}, 
+)

learning3d/losses/emd.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def emd(template: torch.Tensor, source: torch.Tensor):
+	from emd import EMDLoss
+	emd_loss = torch.mean(self.emd(template, source))/(template.size()[1])
+	return emd_loss
+
+
+class EMDLoss(nn.Module):
+	def __init__(self):
+		super(EMDLoss, self).__init__()
+
+	def forward(self, template, source):
+		return emd(template, source)

learning3d/losses/frobenius_norm.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def frobeniusNormLoss(predicted, igt):
+	""" |predicted*igt - I| (should be 0) """
+	assert predicted.size(0) == igt.size(0)
+	assert predicted.size(1) == igt.size(1) and predicted.size(1) == 4
+	assert predicted.size(2) == igt.size(2) and predicted.size(2) == 4
+
+	error = predicted.matmul(igt)
+	I = torch.eye(4).to(error).view(1, 4, 4).expand(error.size(0), 4, 4)
+	return torch.nn.functional.mse_loss(error, I, size_average=True) * 16
+
+
+class FrobeniusNormLoss(nn.Module):
+	def __init__(self):
+		super(FrobeniusNormLoss, self).__init__()
+
+	def forward(self, predicted, igt):
+		return frobeniusNormLoss(predicted, igt)

learning3d/losses/rmse_features.py

@@ -0,0 +1,16 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+def rmseOnFeatures(feature_difference):
+	# |feature_difference| should be 0
+	gt = torch.zeros_like(feature_difference)
+	return torch.nn.functional.mse_loss(feature_difference, gt, size_average=False)
+
+
+class RMSEFeaturesLoss(nn.Module):
+	def __init__(self):
+		super(RMSEFeaturesLoss, self).__init__()
+
+	def forward(self, feature_difference):
+		return rmseOnFeatures(feature_difference)

learning3d/models/__init__.py

@@ -0,0 +1,23 @@
+from .pointnet import PointNet
+from .pointconv import create_pointconv
+from .dgcnn import DGCNN
+from .ppfnet import PPFNet
+from .pooling import Pooling
+
+from .classifier import Classifier
+from .segmentation import Segmentation
+
+from .dcp import DCP
+from .prnet import PRNet
+from .pcrnet import iPCRNet
+from .pointnetlk import PointNetLK
+from .rpmnet import RPMNet
+from .pcn import PCN
+from .deepgmr import DeepGMR
+from .masknet import MaskNet
+from .masknet2 import MaskNet2
+
+try:
+	from .flownet3d import FlowNet3D
+except:
+	print("Error raised in pointnet2 module for FlowNet3D Network!\nEither don't use pointnet2_utils or retry it's setup.")

learning3d/models/classifier.py

@@ -0,0 +1,41 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .pooling import Pooling
+
+class Classifier(nn.Module):
+	def __init__(self, feature_model, num_classes=40):
+		super(Classifier, self).__init__()
+		self.feature_model = feature_model
+		self.num_classes = num_classes
+
+		self.linear1 = torch.nn.Linear(self.feature_model.emb_dims, 512)
+		self.bn1 = torch.nn.BatchNorm1d(512)
+		self.dropout1 = torch.nn.Dropout(p=0.7)
+		self.linear2 = torch.nn.Linear(512, 256)
+		self.bn2 = torch.nn.BatchNorm1d(256)
+		self.dropout2 = torch.nn.Dropout(p=0.7)
+		self.linear3 = torch.nn.Linear(256, self.num_classes)
+
+		self.pooling = Pooling('max')
+
+	def forward(self, input_data):
+		output = self.pooling(self.feature_model(input_data))
+		output = F.relu(self.bn1(self.linear1(output)))
+		output = self.dropout1(output)
+		output = F.relu(self.bn2(self.linear2(output)))
+		output = self.dropout2(output)
+		output = self.linear3(output)
+		return output
+
+
+if __name__ == '__main__':
+	from pointnet import PointNet
+	x = torch.rand(10,1024,3)
+
+	pn = PointNet()
+	classifier = Classifier(pn)
+	classes = classifier(x)
+	
+	print('Input Shape: {}\nClassification Output Shape: {}'
+		  .format(x.shape, classes.shape))

learning3d/models/dcp.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .dgcnn import DGCNN
+from .pointnet import PointNet
+from .. ops import transform_functions as transform
+from .. utils import Transformer, SVDHead, Identity
+
+
+class DCP(nn.Module):
+	def __init__(self, feature_model=DGCNN(), cycle=False, pointer_='transformer', head='svd'):
+		super(DCP, self).__init__()
+		self.cycle = cycle
+		self.emb_nn = feature_model
+
+		if pointer_ == 'identity':
+			self.pointer = Identity()
+		elif pointer_ == 'transformer':
+			self.pointer = Transformer(self.emb_nn.emb_dims, n_blocks=1, dropout=0.0, ff_dims=1024, n_heads=4)
+		else:
+			raise Exception("Not implemented")
+
+		if head == 'mlp':
+			self.head = MLPHead(self.emb_nn.emb_dims)
+		elif head == 'svd':
+			self.head = SVDHead(self.emb_nn.emb_dims)
+		else:
+			raise Exception('Not implemented')
+
+	def forward(self, template, source):
+		source_features = self.emb_nn(source)
+		template_features = self.emb_nn(template)
+
+		source_features_p, template_features_p = self.pointer(source_features, template_features)
+
+		source_features = source_features + source_features_p
+		template_features = template_features + template_features_p
+
+		rotation_ab, translation_ab = self.head(source_features, template_features, source, template)
+		if self.cycle:
+			rotation_ba, translation_ba = self.head(template_features, source_features, template, source)
+		else:
+			rotation_ba = rotation_ab.transpose(2, 1).contiguous()
+			translation_ba = -torch.matmul(rotation_ba, translation_ab.unsqueeze(2)).squeeze(2)
+
+		transformed_source = transform.transform_point_cloud(source, rotation_ab, translation_ab)
+
+		result = {'est_R': rotation_ab,
+				  'est_t': translation_ab,
+				  'est_R_': rotation_ba,
+				  'est_t_': translation_ba,
+				  'est_T': transform.convert2transformation(rotation_ab, translation_ab),
+				  'r': template_features - source_features,
+				  'transformed_source': transformed_source}
+		return result
+
+
+class MLPHead(nn.Module):
+    def __init__(self, emb_dims):
+        super(MLPHead, self).__init__()
+        self.emb_dims = emb_dims
+        self.nn = nn.Sequential(nn.Linear(emb_dims * 2, emb_dims // 2),
+                                nn.BatchNorm1d(emb_dims // 2),
+                                nn.ReLU(),
+                                nn.Linear(emb_dims // 2, emb_dims // 4),
+                                nn.BatchNorm1d(emb_dims // 4),
+                                nn.ReLU(),
+                                nn.Linear(emb_dims // 4, emb_dims // 8),
+                                nn.BatchNorm1d(emb_dims // 8),
+                                nn.ReLU())
+        self.proj_rot = nn.Linear(emb_dims // 8, 4)
+        self.proj_trans = nn.Linear(emb_dims // 8, 3)
+
+    def forward(self, *input):
+        src_embedding = input[0]
+        tgt_embedding = input[1]
+        embedding = torch.cat((src_embedding, tgt_embedding), dim=1)
+        embedding = self.nn(embedding.max(dim=-1)[0])
+        rotation = self.proj_rot(embedding)
+        rotation = rotation / torch.norm(rotation, p=2, dim=1, keepdim=True)
+        translation = self.proj_trans(embedding)
+        return quat2mat(rotation), translation
+
+
+if __name__ == '__main__':
+	template, source = torch.rand(10,1024,3), torch.rand(10,1024,3)
+	pn = PointNet()
+
+	# Not Tested Yet.
+	net = DCP(pn)
+	result = net(template, source)
+	import ipdb; ipdb.set_trace()

learning3d/models/deepgmr.py

@@ -0,0 +1,165 @@
+'''
+We thank the author of DeepGMR paper to open-source their code.
+Modified by Vinit Sarode.
+'''
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .. ops import transform_functions as transform
+
+
+def gmm_params(gamma, pts):
+	'''
+	Inputs:
+		gamma: B x N x J
+		pts: B x N x 3
+	'''
+	# pi: B x J
+	pi = gamma.mean(dim=1)
+	Npi = pi * gamma.shape[1]
+	# mu: B x J x 3
+	mu = gamma.transpose(1, 2) @ pts / Npi.unsqueeze(2)
+	# diff: B x N x J x 3
+	diff = pts.unsqueeze(2) - mu.unsqueeze(1)
+	# sigma: B x J x 3 x 3
+	eye = torch.eye(3).unsqueeze(0).unsqueeze(1).to(gamma.device)
+	sigma = (
+		((diff.unsqueeze(3) @ diff.unsqueeze(4)).squeeze() * gamma).sum(dim=1) / Npi
+	).unsqueeze(2).unsqueeze(3) * eye
+	return pi, mu, sigma
+
+
+def gmm_register(pi_s, mu_s, mu_t, sigma_t):
+	'''
+	Inputs:
+		pi: B x J
+		mu: B x J x 3
+		sigma: B x J x 3 x 3
+	'''
+	c_s = pi_s.unsqueeze(1) @ mu_s
+	c_t = pi_s.unsqueeze(1) @ mu_t
+	Ms = torch.sum((pi_s.unsqueeze(2) * (mu_s - c_s)).unsqueeze(3) @
+				   (mu_t - c_t).unsqueeze(2) @ sigma_t.inverse(), dim=1)
+	U, _, V = torch.svd(Ms.cpu())
+	U = U.cuda() if torch.cuda.is_available() else U
+	V = V.cuda() if torch.cuda.is_available() else V
+	S = torch.eye(3).unsqueeze(0).repeat(U.shape[0], 1, 1).to(U.device)
+	S[:, 2, 2] = torch.det(V @ U.transpose(1, 2))
+	R = V @ S @ U.transpose(1, 2)
+	t = c_t.transpose(1, 2) - R @ c_s.transpose(1, 2)
+	bot_row = torch.Tensor([[[0, 0, 0, 1]]]).repeat(R.shape[0], 1, 1).to(R.device)
+	T = torch.cat([torch.cat([R, t], dim=2), bot_row], dim=1)
+	return T
+
+
+class Conv1dBNReLU(nn.Sequential):
+	def __init__(self, in_planes, out_planes):
+		super(Conv1dBNReLU, self).__init__(
+			nn.Conv1d(in_planes, out_planes, kernel_size=1, bias=False),
+			nn.BatchNorm1d(out_planes),
+			nn.ReLU(inplace=True))
+
+
+class FCBNReLU(nn.Sequential):
+	def __init__(self, in_planes, out_planes):
+		super(FCBNReLU, self).__init__(
+			nn.Linear(in_planes, out_planes, bias=False),
+			nn.BatchNorm1d(out_planes),
+			nn.ReLU(inplace=True))
+
+
+class TNet(nn.Module):
+	def __init__(self):
+		super(TNet, self).__init__()
+		self.encoder = nn.Sequential(
+			Conv1dBNReLU(3, 64),
+			Conv1dBNReLU(64, 128),
+			Conv1dBNReLU(128, 256))
+		self.decoder = nn.Sequential(
+			FCBNReLU(256, 128),
+			FCBNReLU(128, 64),
+			nn.Linear(64, 6))
+
+	@staticmethod
+	def f2R(f):
+		r1 = F.normalize(f[:, :3])
+		proj = (r1.unsqueeze(1) @ f[:, 3:].unsqueeze(2)).squeeze(2)
+		r2 = F.normalize(f[:, 3:] - proj * r1)
+		r3 = r1.cross(r2)
+		return torch.stack([r1, r2, r3], dim=2)
+
+	def forward(self, pts):
+		f = self.encoder(pts)
+		f, _ = f.max(dim=2)
+		f = self.decoder(f)
+		R = self.f2R(f)
+		return R @ pts
+
+
+class PointNet(nn.Module):
+	def __init__(self, use_rri, use_tnet=False, nearest_neighbors=20):
+		super(PointNet, self).__init__()
+		self.use_tnet = use_tnet
+		self.tnet = TNet() if self.use_tnet else None
+		d_input = nearest_neighbors * 4 if use_rri else 3
+		self.encoder = nn.Sequential(
+			Conv1dBNReLU(d_input, 64),
+			Conv1dBNReLU(64, 128),
+			Conv1dBNReLU(128, 256),
+			Conv1dBNReLU(256, args.d_model))
+		self.decoder = nn.Sequential(
+			Conv1dBNReLU(args.d_model * 2, 512),
+			Conv1dBNReLU(512, 256),
+			Conv1dBNReLU(256, 128),
+			nn.Conv1d(128, args.n_clusters, kernel_size=1))
+
+	def forward(self, pts):
+		pts = self.tnet(pts) if self.use_tnet else pts
+		f_loc = self.encoder(pts)
+		f_glob, _ = f_loc.max(dim=2)
+		f_glob = f_glob.unsqueeze(2).expand_as(f_loc)
+		y = self.decoder(torch.cat([f_loc, f_glob], dim=1))
+		return y.transpose(1, 2)
+
+
+class DeepGMR(nn.Module):
+	def __init__(self, use_rri=True, feature_model=None, nearest_neighbors=20):
+		super(DeepGMR, self).__init__()
+		self.backbone = feature_model if not None else PointNet(use_rri=use_rri, nearest_neighbors=nearest_neighbors)
+		self.use_rri = use_rri
+
+	def forward(self, template, source):
+		if self.use_rri:
+			self.template = template[..., :3]
+			self.source = source[..., :3]
+			template_features = template[..., 3:].transpose(1, 2)
+			source_features = source[..., 3:].transpose(1, 2)
+		else:
+			self.template = template
+			self.source = source
+			template_features = (template - template.mean(dim=2, keepdim=True)).transpose(1, 2)
+			source_features = (source - source.mean(dim=2, keepdim=True)).transpose(1, 2)
+
+		self.template_gamma = F.softmax(self.backbone(template_features), dim=2)
+		self.template_pi, self.template_mu, self.template_sigma = gmm_params(self.template_gamma, self.template)
+		self.source_gamma = F.softmax(self.backbone(source_features), dim=2)
+		self.source_pi, self.source_mu, self.source_sigma = gmm_params(self.source_gamma, self.source)
+
+		self.est_T_inverse = gmm_register(self.template_pi, self.template_mu, self.source_mu, self.source_sigma)
+		self.est_T = gmm_register(self.source_pi, self.source_mu, self.template_mu, self.template_sigma) # [template = source * est_T]
+		self.igt = igt				# [source = template * igt]
+
+		transformed_source = transform.transform_point_cloud(source, est_T[:, :3, :3], est_T[:, :3, 3])
+
+		result = {'est_R': est_T[:, :3, :3],
+				  'est_t': est_T[:, :3, 3],
+				  'est_R_inverse': est_T_inverse[:, :3, :3],
+				  'est_t_inverese': est_T_inverse[:, :3, 3],
+				  'est_T': est_T,
+				  'est_T_inverse': est_T_inverse,
+				  'r': template_features - source_features,
+				  'transformed_source': transformed_source}
+
+		return result

learning3d/models/dgcnn.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn.functional as F
+
+def knn(x, k):
+	inner = -2 * torch.matmul(x.transpose(2, 1).contiguous(), x)
+	xx = torch.sum(x ** 2, dim=1, keepdim=True)
+	pairwise_distance = -xx - inner - xx.transpose(2, 1).contiguous()
+
+	idx = pairwise_distance.topk(k=k, dim=-1)[1]  # (batch_size, num_points, k)
+	return idx
+
+
+def get_graph_feature(x, k=20):
+	# x = x.squeeze()
+	idx = knn(x, k=k)  # (batch_size, num_points, k)
+	batch_size, num_points, _ = idx.size()
+	
+	device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+	idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
+
+	idx = idx + idx_base
+
+	idx = idx.view(-1)
+
+	_, num_dims, _ = x.size()
+
+	# (batch_size, num_points, num_dims)  -> (batch_size*num_points, num_dims) #   batch_size * num_points * k + range(0, batch_size*num_points)
+	x = x.transpose(2, 1).contiguous()  
+
+	feature = x.view(batch_size * num_points, -1)[idx, :]
+	feature = feature.view(batch_size, num_points, k, num_dims)
+	x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
+
+	feature = torch.cat((feature, x), dim=3).permute(0, 3, 1, 2)
+
+	return feature
+
+
+class DGCNN(torch.nn.Module):
+	def __init__(self, emb_dims=1024, input_shape="bnc"):
+		super(DGCNN, self).__init__()
+		if input_shape not in ["bcn", "bnc"]:
+			raise ValueError("Allowed shapes are 'bcn' (batch * channels * num_in_points), 'bnc' ")
+		self.input_shape = input_shape
+		self.emb_dims = emb_dims
+
+		self.conv1 = torch.nn.Conv2d(6, 64, kernel_size=1, bias=False)
+		self.conv2 = torch.nn.Conv2d(64, 64, kernel_size=1, bias=False)
+		self.conv3 = torch.nn.Conv2d(64, 128, kernel_size=1, bias=False)
+		self.conv4 = torch.nn.Conv2d(128, 256, kernel_size=1, bias=False)
+		self.conv5 = torch.nn.Conv2d(512, emb_dims, kernel_size=1, bias=False)
+		self.bn1 = torch.nn.BatchNorm2d(64)
+		self.bn2 = torch.nn.BatchNorm2d(64)
+		self.bn3 = torch.nn.BatchNorm2d(128)
+		self.bn4 = torch.nn.BatchNorm2d(256)
+		self.bn5 = torch.nn.BatchNorm2d(emb_dims)
+
+	def forward(self, input_data):
+		if self.input_shape == "bnc":
+			input_data = input_data.permute(0, 2, 1)
+		if input_data.shape[1] != 3:
+			raise RuntimeError("shape of x must be of [Batch x 3 x NumInPoints]")
+
+		batch_size, num_dims, num_points = input_data.size()
+		output = get_graph_feature(input_data)
+
+		output = F.relu(self.bn1(self.conv1(output)))
+		output1 = output.max(dim=-1, keepdim=True)[0]
+
+		output = F.relu(self.bn2(self.conv2(output)))
+		output2 = output.max(dim=-1, keepdim=True)[0]
+
+		output = F.relu(self.bn3(self.conv3(output)))
+		output3 = output.max(dim=-1, keepdim=True)[0]
+
+		output = F.relu(self.bn4(self.conv4(output)))
+		output4 = output.max(dim=-1, keepdim=True)[0]
+
+		output = torch.cat((output1, output2, output3, output4), dim=1)
+
+		output = F.relu(self.bn5(self.conv5(output))).view(batch_size, -1, num_points)
+		return output
+
+
+if __name__ == '__main__':
+	# Test the code.
+	x = torch.rand((10,1024,3))
+
+	dgcnn = DGCNN()
+	y = dgcnn(x)
+	print("\nInput Shape of DGCNN: ", x.shape, "\nOutput Shape of DGCNN: ", y.shape)