def __init__(self, args):
super(ReconstructionNet, self).__init__()
self.encoder = DGCNN_Seg_Encoder(args)
self.decoder = FoldNet_Decoder(args)
if args.loss == 'ChamferLoss':
self.loss = ChamferLoss()
elif args.loss == 'ChamferLoss_m':
self.loss = ChamferLoss_m()
def __init__(self, args):
super(ReconstructionNet, self).__init__()
if args.encoder == 'foldnet':
self.encoder = FoldNet_Encoder(args)
elif args.encoder == 'dgcnn_cls':
self.encoder = DGCNN_Cls_Encoder(args)
elif args.encoder == 'dgcnn_seg':
self.encoder = DGCNN_Seg_Encoder(args)
self.decoder = FoldNet_Decoder(args)
self.loss = ChamferLoss()
def __init__(self, args):
super(MultiTaskNet, self).__init__()
self.encoder = DGCNN_Seg_Encoder(args)
self.decoder = FoldNet_Decoder(args)
if args.rec_loss == 'ChamferLoss':
self.rec_loss = ChamferLoss()
elif args.rec_loss == 'ChamferLoss_m':
self.rec_loss = ChamferLoss_m()
self.classifer = DGCNN_RT_Cls_Classifier(args)
self.rot_loss = CrossEntropyLoss()
def __init__(self, args):
super(DGCNN_FoldNet, self).__init__()
if args.encoder == 'foldingnet':
self.encoder = FoldNet_Encoder(args)
elif args.encoder == 'dgcnn_classification':
self.encoder = DGCNN_Cls_Encoder(args)
elif args.encoder == 'dgcnn_segmentation':
self.encoder = DGCNN_Seg_Encoder(args)
elif args.encoder == 'pointnet':
self.encoder = PointNetEncoder(args)
self.decoder = FoldNet_Decoder(args)
if args.loss == 'ChamferLoss':
self.loss = ChamferLoss()
elif args.loss == 'ChamferLoss_m':
self.loss = ChamferLoss_m()
def __init__(self, num_patches=32, num_points_per_patch=1024):
super(PPFFoldNet, self).__init__()
self.encoder = Encoder(num_patches=num_patches,
num_points_per_patch=num_points_per_patch)
self.decoder = Decoder(num_points_per_patch=num_points_per_patch)
self.loss = ChamferLoss()
# Print the params size of this model.
if torch.cuda.is_available():
summary(self.cuda(), (1024, 4), batch_size=num_patches)
else:
summary(self, (1024, 4), batch_size=num_patches)
for m in self.modules():
if isinstance(m, (nn.Conv1d, nn.Linear)):
nn.init.xavier_normal_(m.weight, gain=1)
elif isinstance(m, (nn.BatchNorm1d)):
nn.init.constant_(m.weight, 1)
nn.init.constant_(m.bias, 0)
torch.backends.cudnn.deterministic = DEBUG
torch.backends.cudnn.benchmark = not DEBUG
print(f'Using device {DEVICE}')
# Data loaders
train_ds = PPFDataset(**TRAIN_DS_ARGS)
train_dl = DataLoader(train_ds, **TRAIN_DL_ARGS)
val_ds = PPFDataset(**VAL_DS_ARGS)
val_dl = DataLoader(val_ds, **VAL_DL_ARGS)
print('Training set: {} Validation set: {}\n'.format(
train_ds.__len__(), val_ds.__len__()
))
# Model
model = AutoEncoder(NUM_PTS_PER_PATCH)
model.apply(init_weights).to(DEVICE)
loss_func = ChamferLoss()
optimizer = Adam(model.parameters(), LR)
scheduler = OneCycleLR(
optimizer, MAX_LR, total_steps=len(train_dl)*TRAINER_ARGS.num_epochs
)
Path(TRAINER_ARGS.checkpoint_path).parent.mkdir(parents=True, exist_ok=True)
# Training
train(model, loss_func, optimizer, scheduler, (train_dl, val_dl), TRAINER_ARGS)
def __init__(self, num_points):
super(FoldNet, self).__init__()
self.encoder = Encoder(num_points=num_points)
self.decoder = Decoder(num_points=num_points)
self.loss = ChamferLoss()