utils.py

import os
import shutil
import time
import pprint

import torch


def set_gpu(x):
    os.environ['CUDA_VISIBLE_DEVICES'] = x
    print('using gpu:', x)


def ensure_path(path):
    if os.path.exists(path):
        if input('{} exists, remove? ([y]/n)'.format(path)) != 'n':
            shutil.rmtree(path)
            os.mkdir(path)
    else:
        os.mkdir(path)

import scipy as sp
import scipy.stats
import numpy as np


class Averager():

    def __init__(self):
        self.n = 0
        self.v = 0
        self.acc = []

    def add(self, x):
        self.v = (self.v * self.n + x) / (self.n + 1)
        self.n += 1

        self.acc.append(x)

    def item(self):
        # return self.v
        a = 1.0 * np.array(self.acc)
        m = np.mean(a)
        return m

    def stat(self, confidence=0.95):
        a = 1.0 * np.array(self.acc)
        assert len(a) == self.n
        se = scipy.stats.sem(a)
        h = se * sp.stats.t._ppf((1 + confidence) / 2., self.n - 1)
        return h


def count_acc(logits, label):
    pred = torch.argmax(logits, dim=1)
    return (pred == label).type(torch.cuda.FloatTensor).mean().item()


def dot_metric(a, b):
    return torch.mm(a, b.t())


def euclidean_metric(a, b):
    n = a.shape[0]
    m = b.shape[0]
    a = a.unsqueeze(1).expand(n, m, -1)
    b = b.unsqueeze(0).expand(n, m, -1)
    logits = -((a - b)**2).sum(dim=2)
    return logits


class Timer():

    def __init__(self):
        self.o = time.time()

    def measure(self, p=1):
        x = (time.time() - self.o) / p
        x = int(x)
        if x >= 3600:
            return '{:.1f}h'.format(x / 3600)
        if x >= 60:
            return '{}m'.format(round(x / 60))
        return '{}s'.format(x)

_utils_pp = pprint.PrettyPrinter()


def pprint(x):
    _utils_pp.pprint(x)


def l2_loss(pred, label):
    return ((pred - label)**2).sum() / len(pred) / 2