class ChamfersDistance(nn.Module): ''' Use NNDModule as a member. ''' def __init__(self): super(ChamfersDistance, self).__init__() self.nnd = NNDModule() def forward(self, input1, input2): # BxNxK, BxMxK dist0, dist1 = self.nnd.forward(input1, input2) # BxN, BxM loss = torch.mean(torch.sqrt(dist0), 1) + torch.mean( torch.sqrt(dist1), 1) # B loss = torch.mean(loss) # 1 # loss = torch.rand(1,requires_grad=True) return loss
def get_chamfer_dist(get_slow=False): try: if get_slow: raise ValueError import sys sys.path.insert(0, './nndistance') from modules.nnd import NNDModule dist = NNDModule() except: dist = chamfer_quadratic def loss(a, b): if a.dim() == 4: if a.size(1) == 2: a = from_polar(a) assert a.size(1) == 3 a = a.permute(0, 2, 3, 1).contiguous().reshape(a.size(0), -1, 3) if b.dim() == 4: if b.size(1) == 2: b = from_polar(b) assert b.size(1) == 3 b = b.permute(0, 2, 3, 1).contiguous().reshape(b.size(0), -1, 3) assert a.dim() == b.dim() == 3 if a.size(-1) != 3: assert a.size(-2) == 3 a = a.transpose(-2, -1).contiguous() if b.size(-1) != 3: assert b.size(-2) == 3 b = a.transpose(-2, -1).contiguous() dist_a, dist_b = dist(a, b) return dist_a.sum(dim=-1) + dist_b.sum(dim=-1) return loss
from ply import * import torch.nn.functional as F import sys from tqdm import tqdm import os import json import time, datetime import subprocess import pandas as pd try: from script.normalize_obj import * except: print('couldnt load normalize obj') sys.path.append("./nndistance/") from modules.nnd import NNDModule distChamfer = NNDModule() parser = argparse.ArgumentParser() parser.add_argument('--batchSize', type=int, default=1, help='input batch size') parser.add_argument('--workers', type=int, help='number of data loading workers', default=6) parser.add_argument('--nepoch', type=int, default=50000, help='number of epochs to train for') parser.add_argument('--model', type=str, default = 'trained_models/ae_baseline.pth', help='yuor path to the trained model') parser.add_argument('--num_points', type=int, default = 2500, help='number of points fed to poitnet') opt = parser.parse_args() print (opt) blue = lambda x:'\033[94m' + x + '\033[0m' opt.manualSeed = random.randint(1, 10000) # fix seed
import torch import torch.nn as nn from torch.autograd import Variable from modules.nnd import NNDModule dist = NNDModule() p1 = torch.rand(1, 3, 3) * 20 p2 = torch.rand(1, 3, 3) * 20 # p1 = p1.int() # p1.random_(0,2) # p1 = p1.float() # p2 = p2.int() # p2.random_(0,2) p2 = p2.float() print(p1) print(p2) print('cpu') points1 = Variable(p1, requires_grad=True) points2 = Variable(p2) idx1, idx2 = dist(points1, points2) print(idx1, idx1) print('gpu') points1_cuda = Variable(p1.cuda(), requires_grad=True) points2_cuda = Variable(p2.cuda()) idx1, idx2 = dist(points1_cuda, points2_cuda) print(idx1, idx2)
def __init__(self): super(ChamfersDistance, self).__init__() self.nnd = NNDModule()
def weights_init(m): classname = m.__class__.__name__ if classname.find('Conv') != -1: m.weight.data.normal_(0.0, 0.02) elif classname.find('BatchNorm') != -1: m.weight.data.normal_(1.0, 0.02) m.bias.data.fill_(0) ae.apply(weights_init) ae.cuda() optimizer = optim.Adagrad(ae.parameters(), lr=0.001) nnd = NNDModule() num_batch = len(dataset) / opt.batchSize for epoch in range(opt.nepoch): for i, data in enumerate(dataloader, 0): optimizer.zero_grad() points, _ = data points = Variable(points) bs = points.size()[0] points = points.transpose(2, 1) points = points.cuda() gen = ae(points)
def __init__(self, n_samples=1024, cuda_opt=True): super(ChamferLoss, self).__init__() self.n_samples = n_samples self.dist = NNDModule() self.cuda_opt = cuda_opt