'--save_model', type=str, default='./save_model/'
)
parser.add_argument(
'--base_model', type=str, default='pointnet'
)
cfg = parser.parse_args()
print(cfg)
model_dict = {
'pointnet': (PointNet(k=40), F.nll_loss, accuracy, torch.optim.Adam, PointNetTrainer),
}
train_transforms = transforms.Compose([
data_utils.PointcloudToTensor(),
data_utils.PointcloudRotatebyAngle(np.pi/4),
data_utils.PointcloudJitter(),
data_utils.PointcloudScaleAndTranslate(),
data_utils.PointcloudScale(),
data_utils.PointcloudTranslate(),
data_utils.PointcloudRandomInputDropout()
])
test_transforms = transforms.Compose([
data_utils.PointcloudToTensor()
])
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
parser.add_argument('--dataset', type=str, default='data/')
parser.add_argument('--workers', type=int, default=4)
parser.add_argument('--save_model', type=str, default='./save_model/')
parser.add_argument('--base_model', type=str, default='pointnet')
cfg = parser.parse_args()
print(cfg)
model_dict = {
'pointnet':
(PointNet(k=40), F.nll_loss, accuracy, torch.optim.Adam, PointNetTrainer)
}
train_transforms = transforms.Compose([
data_utils.PointcloudToTensor(),
data_utils.PointcloudRotatebyAngle(np.pi / 4),
data_utils.PointcloudJitter(),
data_utils.PointcloudScaleAndTranslate(),
data_utils.PointcloudScale(),
data_utils.PointcloudTranslate(),
data_utils.PointcloudRandomInputDropout()
])
test_transforms = transforms.Compose([data_utils.PointcloudToTensor()])
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print(device)
if __name__ == "__main__":
ds_train = ModelNet40Dataset(num_points=2500,