-
Notifications
You must be signed in to change notification settings - Fork 1
/
main_pointcloud.py
93 lines (75 loc) · 3.13 KB
/
main_pointcloud.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import argparse
import torch
from torch_geometric.datasets import ModelNet
from torch_geometric.loader import DataLoader
import torch_geometric.transforms as T
from pytorch_lightning import Trainer, LightningDataModule, seed_everything
from reynet.models.pointcloud import ReyNetModel, DeepsetsModel
from reynet.utils import args_print
class ModelNetDataModule(LightningDataModule):
def __init__(self, name: str, batch_size: int):
super().__init__()
self.save_hyperparameters()
def setup(self, stage):
train_transform = T.Compose([
T.RandomScale((0.8, 1.25)),
T.RandomRotate(18, axis=2),
T.SamplePoints(100),
])
self.trainset = ModelNet(f"./data/ModelNet{self.hparams.name}",
name=self.hparams.name,
train=True,
transform=train_transform,
pre_transform=T.NormalizeScale())
self.testset = ModelNet(f"./data/ModelNet{self.hparams.name}",
name=self.hparams.name,
train=False,
transform=T.SamplePoints(100),
pre_transform=T.NormalizeScale())
def train_dataloader(self):
trainloader = DataLoader(self.trainset,
batch_size=self.hparams.batch_size,
shuffle=True,
num_workers=8)
return trainloader
def val_dataloader(self):
testloader = DataLoader(self.testset,
batch_size=self.hparams.batch_size,
shuffle=False,
num_workers=8)
return testloader
def test_dataloader(self):
testloader = DataLoader(self.testset,
batch_size=self.hparams.batch_size,
shuffle=False,
num_workers=8)
return testloader
def main():
seed_everything(1234)
# Loading Dataset
dm = ModelNetDataModule(args.name, args.batch_size)
if args.model == 'reynet':
model = Model(args)
elif args.model == 'deepset':
model = Model(args.channel_layers, args)
print(model)
trainer = Trainer.from_argparse_args(args, gradient_clip_val=5)
trainer.fit(model, datamodule=dm)
trainer.test(datamodule=dm, ckpt_path='best')
if __name__ == '__main__':
torch.multiprocessing.set_start_method('spawn')
parser = argparse.ArgumentParser()
parser.add_argument('--model', type=str, default='reynet')
temp_args, _ = parser.parse_known_args()
if temp_args.model == 'reynet':
Model = ReyNetModel
elif temp_args.model == 'deepset':
Model = DeepsetsModel
parser.add_argument('--name', type=str, default="40")
parser.add_argument('--batch-size', type=int, default=64)
parser = Model.add_model_specific_args(parser)
parser = Trainer.add_argparse_args(parser)
args = parser.parse_args()
args.max_epochs = 100
args_print(args)
main()