-
Notifications
You must be signed in to change notification settings - Fork 25
/
train.py
44 lines (35 loc) · 1.93 KB
/
train.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
import models.local_model as model
import models.data.voxelized_data_shapenet as voxelized_data
from models import training
import torch
import configs.config_loader as cfg_loader
cfg = cfg_loader.get_config()
net = model.NDF()
train_dataset = voxelized_data.VoxelizedDataset('train',
res=cfg.input_res,
pointcloud_samples=cfg.num_points,
data_path=cfg.data_dir,
split_file=cfg.split_file,
batch_size=cfg.batch_size,
num_sample_points=cfg.num_sample_points_training,
num_workers=30,
sample_distribution=cfg.sample_ratio,
sample_sigmas=cfg.sample_std_dev)
val_dataset = voxelized_data.VoxelizedDataset('val',
res=cfg.input_res,
pointcloud_samples=cfg.num_points,
data_path=cfg.data_dir,
split_file=cfg.split_file,
batch_size=cfg.batch_size,
num_sample_points=cfg.num_sample_points_training,
num_workers=30,
sample_distribution=cfg.sample_ratio,
sample_sigmas=cfg.sample_std_dev)
trainer = training.Trainer(net,
torch.device("cuda"),
train_dataset,
val_dataset,
cfg.exp_name,
optimizer=cfg.optimizer,
lr=cfg.lr)
trainer.train_model(cfg.num_epochs)