utils.py

import torch.nn as nn
import os
import torch
import matplotlib.pyplot as plt
import warnings


def count_params(model):
    tot_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print("Tot params: %d - Trainable params: %d" % (tot_params, trainable_params))
    return tot_params, trainable_params


def plot_grad_flow(named_parameters):
    """
        Credit to RoshanRane: https://discuss.pytorch.org/t/check-gradient-flow-in-network/15063/8
    """
    ave_grads = []
    layers = []
    for n, p in named_parameters:
        if (p.requires_grad) and ("bias" not in n):
            if p.grad is None:
                warnings.warn('Grad is none for layer {}.'.format(n), RuntimeWarning)
                continue
            layers.append(n)
            ave_grads.append(p.grad.abs().mean())
    plt.plot(ave_grads, alpha=0.3, color="b")
    plt.hlines(0, 0, len(ave_grads) + 1, linewidth=1, color="k")
    plt.xticks(range(0, len(ave_grads), 1), layers, rotation="vertical")
    plt.xlim(xmin=0, xmax=len(ave_grads))
    plt.xlabel("Layers")
    plt.ylabel("Average gradient")
    plt.title("Gradient flow")
    plt.grid(True)
    plt.tight_layout()
    return plt.gcf()


class IOStream:
    def __init__(self, path):
        self.f = open(path, 'a')
        self.blue = lambda x: '\033[94m' + x + '\033[0m'
        self.red = lambda x: '\033[31m' + x + '\033[0m'

    def cprint(self, text, color=None):
        if color is not None and (color == 'b' or color == 'blue'):
            print(self.blue(text))
        elif color is not None and (color == 'r' or color == 'red'):
            print(self.red(text))
        else:
            print(text)

        self.f.write(text + '\n')
        self.f.flush()

    def fprint(self, text):
        self.f.write(text + '\n')
        self.f.flush()

    def close(self):
        self.f.close()


def safe_make_dirs(dirs):
    for currdir in dirs:
        if not os.path.exists(currdir):
            os.makedirs(currdir)

###################
#   PFNet stuff   #
###################


def array2samples_distance(array1, array2):
    """
    arguments: 
        array1: the array, size: (num_point, num_feature)
        array2: the samples, size: (num_point, num_feature)
    returns:
        distances: each entry is the distance from a sample to array1 
    """
    num_point1, num_features1 = array1.shape
    num_point2, num_features2 = array2.shape
    expanded_array1 = array1.repeat(num_point2, 1)
    expanded_array2 = torch.reshape(
        torch.unsqueeze(array2, 1).repeat(1, num_point1, 1),
        (-1, num_features2))
    distances = (expanded_array1 - expanded_array2) * (expanded_array1 - expanded_array2)
    #    distances = torch.sqrt(distances)
    distances = torch.sum(distances, dim=1)
    distances = torch.reshape(distances, (num_point2, num_point1))
    distances = torch.min(distances, dim=1)[0]
    distances = torch.mean(distances)
    return distances


def chamfer_distance_numpy(array1, array2):
    batch_size, num_point, num_features = array1.shape
    dist = 0
    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])
        dist = dist + (0.5 * av_dist1 + 0.5 * av_dist2) / batch_size
    return dist * 100


def chamfer_distance_numpy_test(array1, array2):
    batch_size, num_point, num_features = array1.shape
    dist_all = 0
    dist1 = 0
    dist2 = 0
    for i in range(batch_size):
        av_dist1 = array2samples_distance(array1[i], array2[i])
        av_dist2 = array2samples_distance(array2[i], array1[i])
        dist_all = dist_all + (av_dist1 + av_dist2) / batch_size
        dist1 = dist1 + av_dist1 / batch_size
        dist2 = dist2 + av_dist2 / batch_size
    return dist_all, dist1, dist2


class PointLoss(nn.Module):
    def __init__(self):
        super(PointLoss, self).__init__()

    def forward(self, array1, array2):
        return chamfer_distance_numpy(array1, array2)


class PointLoss_test(nn.Module):
    def __init__(self):
        super(PointLoss_test, self).__init__()

    def forward(self, array1, array2):
        return chamfer_distance_numpy_test(array1, array2)


def distance_squre(p1, p2):
    tensor = torch.FloatTensor(p1) - torch.FloatTensor(p2)
    val = tensor.mul(tensor)
    val = val.sum()
    return val


def index_points(points, idx):
    """
    Input:
        points: input points data, [B, N, C]
        idx: sample index data, [B, S]
    Return:
        new_points:, indexed points data, [B, S, C]
    """
    device = points.device
    B = points.shape[0]
    view_shape = list(idx.shape)
    view_shape[1:] = [1] * (len(view_shape) - 1)
    repeat_shape = list(idx.shape)
    repeat_shape[0] = 1
    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
    new_points = points[batch_indices, idx, :]
    return new_points


def farthest_point_sample(xyz, npoint, RAN=True):
    """
    Input:
        xyz: pointcloud data, [B, N, C]
        npoint: number of samples
    Return:
        centroids: sampled pointcloud index, [B, npoint]
    """
    device = xyz.device
    B, N, C = xyz.shape
    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
    distance = torch.ones(B, N).to(device) * 1e10
    if RAN:
        farthest = torch.randint(0, 1, (B,), dtype=torch.long).to(device)
    else:
        farthest = torch.randint(1, 2, (B,), dtype=torch.long).to(device)

    batch_indices = torch.arange(B, dtype=torch.long).to(device)
    for i in range(npoint):
        centroids[:, i] = farthest
        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
        dist = torch.sum((xyz - centroid) ** 2, -1)
        mask = dist < distance
        distance[mask] = dist[mask]
        farthest = torch.max(distance, -1)[1]
    return centroids