def create_transform(config):
if config["model"]["name"] == "pointnet":
transform = T.ToTensor()
elif config["model"]["name"] == "pointnet2":
transform = T.ToTensor()
elif config["model"]["name"] == "dgcnn":
transform = T.ToTensor()
return transform
def test_compose():
size = torch.Size([3, 1024])
x = np.random.rand(2048, 3)
t = T.Compose([T.Shuffle(), T.RandomPointSample(1024), T.ToTensor()])
x = t(x)
assert x.is_contiguous()
assert x.shape == size
def main(args):
# ToTensor transform converts numpy point clouds into tensors
transform = transforms.ToTensor()
dataloaders = {
"train":
data.DataLoader(
dsets.ModelNet40(
args.root,
train=True,
transform=transform,
download=True, # download dataset if needed
),
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=True,
),
"test":
data.DataLoader(
dsets.ModelNet40(
args.root,
train=False, # now we use the test set
transform=transform,
download=False,
),
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=False,
),
}
# Create PointNet model and its optimizer
model = models.PointNet(args.in_channels, args.num_classes).to(args.device)
optimizer = optim.Adam(model.parameters(), args.lr)
criteria = nn.CrossEntropyLoss()
# Showing random samples
points, labels = next(iter(dataloaders["test"]))
points = points.transpose(2, 1).numpy()
visualize(points, labels)
# Here comes the training loop
for epoch in range(args.epochs):
train_epoch(args, epoch, model, dataloaders["train"], optimizer,
criteria)
predict = evaluate(args, model, dataloaders["test"])
visualize(points, labels, predict)
print("Done.")
def main(args):
# ToTensor transform converts numpy point clouds into tensors
transform = transforms.ToTensor()
dataloaders = {
"train":
data.DataLoader(
dsets.ShapeNetPart(
args.root,
split="train",
transform=transform,
download=True, # download dataset if needed
),
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=True,
),
"test":
data.DataLoader(
dsets.ShapeNetPart(
args.root,
split="test", # now we use the test set
transform=transform,
download=False,
),
batch_size=args.batch_size,
num_workers=args.num_workers,
pin_memory=True,
shuffle=False,
),
}
# Create FoldingNet model
model = FoldingNet(args.in_channels).to(args.device)
optimizer = optim.Adam(model.parameters(), args.lr)
criteria = ChamferLoss()
points, _ = next(iter(dataloaders["test"]))
points = points.transpose(2, 1).numpy()
# Here comes the training loop
for epoch in range(args.epochs):
train_epoch(args, epoch, model, dataloaders["train"], optimizer,
criteria)
predict = evaluate(args, model, dataloaders["test"])
visualize(points, predict[0])
def test_to_tensor():
size = torch.Size([3, 2048])
x = np.random.rand(2048, 3)
x = T.ToTensor()(x)
assert x.is_contiguous()
assert x.shape == size