learning3d 0.0.1__py3-none-any.whl

learning3d/examples/train_flownet.py

@@ -0,0 +1,259 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+
+from __future__ import print_function
+import os
+import gc
+import argparse
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+from torch.optim.lr_scheduler import MultiStepLR
+from learning3d.models import FlowNet3D
+from learning3d.data_utils import SceneflowDataset
+import numpy as np
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+class IOStream:
+    def __init__(self, path):
+        self.f = open(path, 'a')
+
+    def cprint(self, text):
+        print(text)
+        self.f.write(text + '\n')
+        self.f.flush()
+
+    def close(self):
+        self.f.close()
+
+
+def _init_(args):
+    if not os.path.exists('checkpoints'):
+        os.makedirs('checkpoints')
+    if not os.path.exists('checkpoints/' + args.exp_name):
+        os.makedirs('checkpoints/' + args.exp_name)
+    if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
+        os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
+
+def weights_init(m):
+    classname=m.__class__.__name__
+    if classname.find('Conv2d') != -1:
+        nn.init.kaiming_normal_(m.weight.data)
+    if classname.find('Conv1d') != -1:
+        nn.init.kaiming_normal_(m.weight.data)
+
+def test_one_epoch(args, net, test_loader):
+    net.eval()
+
+    total_loss = 0
+    num_examples = 0
+    for i, data in tqdm(enumerate(test_loader), total=len(test_loader), smoothing=0.9):
+        pc1, pc2, color1, color2, flow, mask1 = data
+        pc1 = pc1.cuda().transpose(2,1).contiguous()
+        pc2 = pc2.cuda().transpose(2,1).contiguous()
+        color1 = color1.cuda().transpose(2,1).contiguous()
+        color2 = color2.cuda().transpose(2,1).contiguous()
+        flow = flow.cuda()
+        mask1 = mask1.cuda().float()
+
+        batch_size = pc1.size(0)
+        num_examples += batch_size
+        flow_pred = net(pc1, pc2, color1, color2).permute(0,2,1)
+        loss_1 = torch.mean(mask1 * torch.sum((flow_pred - flow) * (flow_pred - flow), -1) / 2.0)
+
+        pc1, pc2 = pc1.permute(0,2,1), pc2.permute(0,2,1)
+        pc1_ = pc1 + flow_pred
+
+        total_loss += loss_1.item() * batch_size
+        
+
+    return total_loss * 1.0 / num_examples
+
+
+def train_one_epoch(args, net, train_loader, opt):
+    net.train()
+    num_examples = 0
+    total_loss = 0
+    for i, data in tqdm(enumerate(train_loader), total=len(train_loader), smoothing=0.9):
+        pc1, pc2, color1, color2, flow, mask1 = data
+        pc1 = pc1.cuda().transpose(2,1).contiguous()
+        pc2 = pc2.cuda().transpose(2,1).contiguous()
+        color1 = color1.cuda().transpose(2,1).contiguous()
+        color2 = color2.cuda().transpose(2,1).contiguous()
+        flow = flow.cuda().transpose(2,1).contiguous()
+        mask1 = mask1.cuda().float()
+
+        batch_size = pc1.size(0)
+        opt.zero_grad()
+        num_examples += batch_size
+        flow_pred = net(pc1, pc2, color1, color2)
+        loss_1 = torch.mean(mask1 * torch.sum((flow_pred - flow) ** 2, 1) / 2.0)
+
+        pc1, pc2, flow_pred = pc1.permute(0,2,1), pc2.permute(0,2,1), flow_pred.permute(0,2,1)
+        pc1_ = pc1 + flow_pred
+
+        loss_1.backward()
+
+        opt.step()
+        total_loss += loss_1.item() * batch_size
+
+        # if (i+1) % 100 == 0:
+        #     print("batch: %d, mean loss: %f" % (i, total_loss / 100 / batch_size))
+        #     total_loss = 0
+    return total_loss * 1.0 / num_examples
+
+
+def test(args, net, test_loader, boardio, textio):
+
+    test_loss = test_one_epoch(args, net, test_loader)
+
+    textio.cprint('==FINAL TEST==')
+    textio.cprint('mean test loss: %f'%test_loss)
+
+
+def train(args, net, train_loader, test_loader, boardio, textio):
+    if args.use_sgd:
+        print("Use SGD")
+        opt = optim.SGD(net.parameters(), lr=args.lr * 100, momentum=args.momentum, weight_decay=1e-4)
+    else:
+        print("Use Adam")
+        opt = optim.Adam(net.parameters(), lr=args.lr, weight_decay=1e-4)
+    scheduler = MultiStepLR(opt, milestones=[75, 150, 200], gamma=0.1)
+
+    best_test_loss = np.inf
+    for epoch in range(args.epochs):
+        scheduler.step()
+        textio.cprint('==epoch: %d=='%epoch)
+        train_loss = train_one_epoch(args, net, train_loader, opt)
+        textio.cprint('mean train EPE loss: %f'%train_loss)
+
+        test_loss = test_one_epoch(args, net, test_loader)
+        textio.cprint('mean test EPE loss: %f'%test_loss)
+
+        if best_test_loss >= test_loss:
+            best_test_loss = test_loss
+            textio.cprint('best test loss till now: %f'%test_loss)
+            if torch.cuda.device_count() > 1:
+                torch.save(net.module.state_dict(), 'checkpoints/%s/models/model.best.t7' % args.exp_name)
+            else:
+                torch.save(net.state_dict(), 'checkpoints/%s/models/model.best.t7' % args.exp_name)
+
+        boardio.add_scalar('Train Loss', train_loss, epoch+1)
+        boardio.add_scalar('Test Loss', test_loss, epoch+1)
+        boardio.add_scalar('Best Test Loss', best_test_loss, epoch+1)
+        
+        if torch.cuda.device_count() > 1:
+            torch.save(net.module.state_dict(), 'checkpoints/%s/models/model.%d.t7' % (args.exp_name, epoch))
+        else:
+            torch.save(net.state_dict(), 'checkpoints/%s/models/model.%d.t7' % (args.exp_name, epoch))
+        gc.collect()
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Point Cloud Registration')
+    parser.add_argument('--exp_name', type=str, default='exp_flownet', metavar='N',
+                        help='Name of the experiment')
+    parser.add_argument('--model', type=str, default='flownet', metavar='N',
+                        choices=['flownet'],
+                        help='Model to use, [flownet]')
+    parser.add_argument('--emb_dims', type=int, default=512, metavar='N',
+                        help='Dimension of embeddings')
+    parser.add_argument('--num_points', type=int, default=2048,
+                        help='Point Number [default: 2048]')
+    parser.add_argument('--dropout', type=float, default=0.5, metavar='N',
+                        help='Dropout ratio in transformer')
+    parser.add_argument('--batch_size', type=int, default=16, metavar='batch_size',
+                        help='Size of batch)')
+    parser.add_argument('--test_batch_size', type=int, default=10, metavar='batch_size',
+                        help='Size of batch)')
+    parser.add_argument('--epochs', type=int, default=250, metavar='N',
+                        help='number of episode to train ')
+    parser.add_argument('--use_sgd', action='store_true', default=True,
+                        help='Use SGD')
+    parser.add_argument('--lr', type=float, default=0.001, metavar='LR',
+                        help='learning rate (default: 0.001, 0.1 if using sgd)')
+    parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
+                        help='SGD momentum (default: 0.9)')
+    parser.add_argument('--no_cuda', action='store_true', default=False,
+                        help='enables CUDA training')
+    parser.add_argument('--seed', type=int, default=1234, metavar='S',
+                        help='random seed (default: 1)')
+    parser.add_argument('--eval', action='store_true', default=False,
+                        help='evaluate the model')
+    parser.add_argument('--cycle', type=bool, default=False, metavar='N',
+                        help='Whether to use cycle consistency')
+    parser.add_argument('--gaussian_noise', type=bool, default=False, metavar='N',
+                        help='Wheter to add gaussian noise')
+    parser.add_argument('--unseen', type=bool, default=False, metavar='N',
+                        help='Whether to test on unseen category')
+    parser.add_argument('--dataset', type=str, default='SceneflowDataset',
+                        choices=['SceneflowDataset'], metavar='N',
+                        help='dataset to use')
+    parser.add_argument('--dataset_path', type=str, default='data_processed_maxcut_35_20k_2k_8192', metavar='N',
+                        help='dataset to use')
+    parser.add_argument('--model_path', type=str, default='', metavar='N',
+                        help='Pretrained model path')
+    parser.add_argument('--pretrained', type=str, default='', metavar='N',
+                        help='Pretrained model path')
+
+    args = parser.parse_args()
+    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
+    # CUDA settings
+    torch.backends.cudnn.deterministic = True
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    np.random.seed(args.seed)
+
+    boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
+    _init_(args)
+
+    textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
+    textio.cprint(str(args))
+
+    if args.dataset == 'SceneflowDataset':
+        train_loader = DataLoader(
+            SceneflowDataset(npoints=args.num_points, partition='train'),
+            batch_size=args.batch_size, shuffle=True, drop_last=True)
+        test_loader = DataLoader(
+            SceneflowDataset(npoints=args.num_points, partition='test'),
+            batch_size=args.test_batch_size, shuffle=False, drop_last=False)
+    else:
+        raise Exception("not implemented")
+
+    if args.model == 'flownet':
+        net = FlowNet3D().cuda()
+        net.apply(weights_init)
+        if args.pretrained:
+            net.load_state_dict(torch.load(args.pretrained), strict=False)
+            print("Pretrained Model Loaded Successfully!")
+        if args.eval:
+            if args.model_path is '':
+                model_path = 'checkpoints' + '/' + args.exp_name + '/models/model.best.t7'
+            else:
+                model_path = args.model_path
+                print(model_path)
+            if not os.path.exists(model_path):
+                print("can't find pretrained model")
+                return
+            net.load_state_dict(torch.load(model_path), strict=False)
+        if torch.cuda.device_count() > 1:
+            net = nn.DataParallel(net)
+            print("Let's use", torch.cuda.device_count(), "GPUs!")
+    else:
+        raise Exception('Not implemented')
+    if args.eval:
+        test(args, net, test_loader, boardio, textio)
+    else:
+        train(args, net, train_loader, test_loader, boardio, textio)
+
+
+    print('FINISH')
+    # boardio.close()
+
+
+if __name__ == '__main__':
+    main()

learning3d/examples/train_masknet.py

@@ -0,0 +1,239 @@
+import argparse
+import os
+import sys
+import logging
+import numpy
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from torch.utils.data import DataLoader
+from tensorboardX import SummaryWriter
+from tqdm import tqdm
+
+# Only if the files are in example folder.
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+if BASE_DIR[-8:] == 'examples':
+	sys.path.append(os.path.join(BASE_DIR, os.pardir))
+	os.chdir(os.path.join(BASE_DIR, os.pardir))
+	
+from learning3d.models import MaskNet
+from learning3d.data_utils import RegistrationData, ModelNet40Data
+
+def _init_(args):
+	if not os.path.exists('checkpoints'):
+		os.makedirs('checkpoints')
+	if not os.path.exists('checkpoints/' + args.exp_name):
+		os.makedirs('checkpoints/' + args.exp_name)
+	if not os.path.exists('checkpoints/' + args.exp_name + '/' + 'models'):
+		os.makedirs('checkpoints/' + args.exp_name + '/' + 'models')
+	os.system('cp train.py checkpoints' + '/' + args.exp_name + '/' + 'train.py.backup')
+	os.system('cp learning3d/models/masknet.py checkpoints' + '/' + args.exp_name + '/' + 'masknet.py.backup')
+	os.system('cp learning3d/data_utils/dataloaders.py checkpoints' + '/' + args.exp_name + '/' + 'dataloaders.py.backup')
+
+
+class IOStream:
+	def __init__(self, path):
+		self.f = open(path, 'a')
+
+	def cprint(self, text):
+		print(text)
+		self.f.write(text + '\n')
+		self.f.flush()
+
+	def close(self):
+		self.f.close()
+
+def test_one_epoch(args, model, test_loader):
+	model.eval()
+	test_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(test_loader)):
+		template, source, igt, gt_mask = data
+
+		template = template.to(args.device)
+		source = source.to(args.device)
+		igt = igt.to(args.device)				# [source] = [igt]*[template]
+		gt_mask = gt_mask.to(args.device)
+
+		masked_template, predicted_mask = model(template, source)
+
+		if args.loss_fn == 'mse':
+			loss_mask = torch.nn.functional.mse_loss(predicted_mask, gt_mask)
+		elif args.loss_fn == 'bce':
+			loss_mask = torch.nn.BCELoss()(predicted_mask, gt_mask)
+
+		test_loss += loss_mask.item()
+		count += 1
+
+	test_loss = float(test_loss)/count
+	return test_loss
+
+def test(args, model, test_loader, textio):
+	test_loss, test_accuracy = test_one_epoch(args.device, model, pnlk, test_loader)
+	textio.cprint('Validation Loss: %f & Validation Accuracy: %f'%(test_loss, test_accuracy))
+
+def train_one_epoch(args, model, train_loader, optimizer):
+	model.train()
+	train_loss = 0.0
+	pred  = 0.0
+	count = 0
+	for i, data in enumerate(tqdm(train_loader)):
+		template, source, igt, gt_mask = data
+
+		template = template.to(args.device)
+		source = source.to(args.device)
+		igt = igt.to(args.device)					# [source] = [igt]*[template]
+		gt_mask = gt_mask.to(args.device)
+
+		masked_template, predicted_mask = model(template, source)
+
+		if args.loss_fn == 'mse':
+			loss_mask = torch.nn.functional.mse_loss(predicted_mask, gt_mask)
+		elif args.loss_fn == 'bce':
+			loss_mask = torch.nn.BCELoss()(predicted_mask, gt_mask)
+
+		# forward + backward + optimize
+		optimizer.zero_grad()
+		loss_mask.backward()
+		optimizer.step()
+
+		train_loss += loss_mask.item()
+		count += 1
+
+	train_loss = float(train_loss)/count
+	return train_loss
+
+def train(args, model, train_loader, test_loader, boardio, textio, checkpoint):
+	learnable_params = filter(lambda p: p.requires_grad, model.parameters())
+	if args.optimizer == 'Adam':
+		optimizer = torch.optim.Adam(learnable_params, lr=0.0001)
+	else:
+		optimizer = torch.optim.SGD(learnable_params, lr=0.1)
+
+	if checkpoint is not None:
+		min_loss = checkpoint['min_loss']
+		optimizer.load_state_dict(checkpoint['optimizer'])
+
+	best_test_loss = np.inf
+
+	for epoch in range(args.start_epoch, args.epochs):
+		train_loss = train_one_epoch(args, model, train_loader, optimizer)
+		test_loss = test_one_epoch(args, model, test_loader)
+
+		if test_loss<best_test_loss:
+			best_test_loss = test_loss
+			
+			snap = {'epoch': epoch + 1,
+					'model': model.state_dict(),
+					'min_loss': best_test_loss,
+					'optimizer' : optimizer.state_dict(),}
+			torch.save(snap, 'checkpoints/%s/models/best_model_snap.t7' % (args.exp_name))
+			torch.save(model.state_dict(), 'checkpoints/%s/models/best_model.t7' % (args.exp_name))
+
+		snap = {'epoch': epoch + 1,
+				'model': model.state_dict(),
+				'min_loss': best_test_loss,
+				'optimizer' : optimizer.state_dict(),}
+		torch.save(snap, 'checkpoints/%s/models/model_snap.t7' % (args.exp_name))
+		torch.save(model.state_dict(), 'checkpoints/%s/models/model.t7' % (args.exp_name))
+		
+		boardio.add_scalar('Train_Loss', train_loss, epoch+1)
+		boardio.add_scalar('Test_Loss', test_loss, epoch+1)
+		boardio.add_scalar('Best_Test_Loss', best_test_loss, epoch+1)
+
+		textio.cprint('EPOCH:: %d, Traininig Loss: %f, Testing Loss: %f, Best Loss: %f'%(epoch+1, train_loss, test_loss, best_test_loss))
+
+def options():
+	parser = argparse.ArgumentParser(description='MaskNet: A Fully-Convolutional Network For Inlier Estimation (Training)')
+	parser.add_argument('--exp_name', type=str, default='exp_masknet', metavar='N',
+						help='Name of the experiment')
+	parser.add_argument('--eval', type=bool, default=False, help='Train or Evaluate the network.')
+	
+	# settings for input data
+	parser.add_argument('--num_points', default=1024, type=int,
+						metavar='N', help='points in point-cloud (default: 1024)')
+	parser.add_argument('--partial_source', default=True, type=bool,
+						help='create partial source point cloud in dataset.')
+	parser.add_argument('--noise', default=False, type=bool,
+						help='Add noise in source point clouds.')
+	parser.add_argument('--outliers', default=False, type=bool,
+						help='Add outliers to template point cloud.')
+
+	# settings for on training
+	parser.add_argument('--seed', type=int, default=1234)
+	parser.add_argument('-j', '--workers', default=4, type=int,
+						metavar='N', help='number of data loading workers (default: 4)')
+	parser.add_argument('-b', '--batch_size', default=32, type=int,
+						metavar='N', help='mini-batch size (default: 32)')
+	parser.add_argument('--test_batch_size', default=8, type=int,
+						metavar='N', help='test-mini-batch size (default: 8)')
+	parser.add_argument('--unseen', default=False, type=bool,
+						help='Use first 20 categories for training and last 20 for testing')
+	parser.add_argument('--epochs', default=500, type=int,
+						metavar='N', help='number of total epochs to run')
+	parser.add_argument('--start_epoch', default=0, type=int,
+						metavar='N', help='manual epoch number (useful on restarts)')
+	parser.add_argument('--optimizer', default='Adam', choices=['Adam', 'SGD'],
+						metavar='METHOD', help='name of an optimizer (default: Adam)')
+	parser.add_argument('--resume', default='', type=str,
+						metavar='PATH', help='path to latest checkpoint (default: null (no-use))')
+	parser.add_argument('--pretrained', default='', type=str,
+						metavar='PATH', help='path to pretrained model file (default: null (no-use))')
+	parser.add_argument('--device', default='cuda:0', type=str,
+						metavar='DEVICE', help='use CUDA if available')
+	parser.add_argument('--loss_fn', default='mse', type=str, choices=['mse', 'bce'])
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = options()
+
+	torch.backends.cudnn.deterministic = True
+	torch.manual_seed(args.seed)
+	torch.cuda.manual_seed_all(args.seed)
+	np.random.seed(args.seed)
+
+	boardio = SummaryWriter(log_dir='checkpoints/' + args.exp_name)
+	_init_(args)
+
+	textio = IOStream('checkpoints/' + args.exp_name + '/run.log')
+	textio.cprint(str(args))
+
+	trainset = RegistrationData(ModelNet40Data(train=True, num_points=args.num_points, unseen=args.unseen),
+								partial_source=args.partial_source, noise=args.noise, outliers=args.outliers, 
+								additional_params={'use_masknet': True})
+	testset = RegistrationData(ModelNet40Data(train=False, num_points=args.num_points, unseen=args.unseen),
+								partial_source=args.partial_source, noise=args.noise, outliers=args.outliers,
+								additional_params={'use_masknet': True})
+	train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers)
+	test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=False, drop_last=False, num_workers=args.workers)
+
+	if not torch.cuda.is_available():
+		args.device = 'cpu'
+	args.device = torch.device(args.device)
+
+	model = MaskNet()
+	model = model.to(args.device)
+
+	checkpoint = None
+	if args.resume:
+		assert os.path.isfile(args.resume)
+		checkpoint = torch.load(args.resume)
+		args.start_epoch = checkpoint['epoch']
+		model.load_state_dict(checkpoint['model'])
+
+	if args.pretrained:
+		assert os.path.isfile(args.pretrained)
+		model.load_state_dict(torch.load(args.pretrained, map_location='cpu'))
+	model.to(args.device)
+
+	if args.eval:
+		test(args, model, test_loader, textio)
+	else:
+		train(args, model, train_loader, test_loader, boardio, textio, checkpoint)
+
+if __name__ == '__main__':
+	main()