utils.py

import argparse
import torch
import os
import numpy as np
from torch import nn
from torch.autograd import Variable
from pathlib import Path

class One_Hot(nn.Module):
    # from :
    # https://lirnli.wordpress.com/2017/09/03/one-hot-encoding-in-pytorch/
    def __init__(self, depth):
        super(One_Hot,self).__init__()
        self.depth = depth
        self.ones = torch.eye(depth)
    def forward(self, X_in):
        X_in = X_in.long()
        return Variable(self.ones.index_select(0,X_in.data))
    def __repr__(self):
        return self.__class__.__name__ + "({})".format(self.depth)


def cuda(tensor,is_cuda):
    if is_cuda : return tensor.cuda()
    else : return tensor

def str2bool(v):
    # codes from : https://stackoverflow.com/questions/15008758/parsing-boolean-values-with-argparse

    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')


def print_network(net):
    num_params = 0
    for param in net.parameters():
        num_params += param.numel()
    print(net)
    print('Total number of parameters: %d' % num_params)


def rm_dir(dir_path, silent=True):
    p = Path(dir_path).resolve()
    if (not p.is_file()) and (not p.is_dir()) :
        print('It is not path for file nor directory :',p)
        return

    paths = list(p.iterdir())
    if (len(paths) == 0) and p.is_dir() :
        p.rmdir()
        if not silent : print('removed empty dir :',p)

    else :
        for path in paths :
            if path.is_file() :
                path.unlink()
                if not silent : print('removed file :',path)
            else:
                rm_dir(path)
        p.rmdir()
        if not silent : print('removed empty dir :',p)

def where(cond, x, y):
    """
    code from :
        https://discuss.pytorch.org/t/how-can-i-do-the-operation-the-same-as-np-where/1329/8
    """
    cond = cond.float()
    return (cond*x) + ((1-cond)*y)
def one_hotify(x, n_classes=10):
    x_long = torch.LongTensor(x)
    x_onehot = torch.eye(n_classes).index_select(0, x_long)
    return x_onehot
class Trainer():
    def __init__(self, model, optimizer, criterion,
                 trn_loader, tst_loader,
                 one_hot=False, use_reconstructions=False, use_cuda=False,
                 print_every=None):
        self.model = model
        self.criterion = criterion
        self.optimizer = optimizer
        self.trn_loader = trn_loader
        self.tst_loader = tst_loader
        
        self.one_hot = one_hot
        self.use_reconstructions = use_reconstructions
        self.use_cuda = use_cuda
        
        self.metrics = {
            'loss': {
                'trn':[],
                'tst':[]
            },
            'accuracy': {
                'trn':[],
                'tst':[]
            },
        }
        self.print_every = print_every
    @staticmethod
    def drawProgressBar(percent, barLen = 50):
        sys.stdout.write("\r")
        progress = ""
        for i in range(barLen):
            if i<int(barLen * percent):
                progress += "="
            else:
                progress += " "
        sys.stdout.write("[ %s ] %.2f%%" % (progress, percent * 100))
        sys.stdout.flush()
        
    def run(self, epochs):
        print('[*] Training for %d epochs' % epochs)
        for epoch in range(1, epochs+1):
            trn_loss, trn_acc = self.train()
            tst_loss, tst_acc = self.test()
            print('[*] Epoch %d, TrnLoss: %.3f, TrnAcc: %.3f, TstLoss: %.3f, TstAcc: %.3f'
                % (epoch, trn_loss, trn_acc, tst_loss, tst_acc))
            self.metrics['accuracy']['trn'].append(trn_acc)
            self.metrics['accuracy']['tst'].append(tst_acc)
            self.metrics['loss']['trn'].append(trn_loss)
            self.metrics['loss']['tst'].append(tst_loss)
    
    def train(self):
        self.model.train()
        return self.run_epoch(self.trn_loader)
    
    def test(self):
        self.model.eval()
        return self.run_epoch(self.tst_loader, train=False)
    
    def make_var(self, tensor):
        if self.use_cuda:
            tensor = tensor.cuda()
        return Variable(tensor)
    
    def run_epoch(self, loader, train=True):
        n_batches = len(loader)
        cum_loss = 0
        cum_acc = 0
        for i, (X, y) in enumerate(loader):
            Trainer.drawProgressBar(i/n_batches)
            X_var = self.make_var(X)
            y_var = self.make_var(one_hotify(y) if self.one_hot else y)

            if self.use_reconstructions:
                scores, reconstructions = self.model(X_var, y_var)
                loss_var = self.criterion(scores, y_var, reconstructions, X_var)
            else:
                scores = self.model(X_var)
                loss_var = self.criterion(scores, y_var)

            if train:
                self.optimizer.zero_grad()
                loss_var.backward()
                self.optimizer.step()

            pred = get_argmax(scores)
            #acc = get_accuracy(pred, y)
            acc = torch.eq(pred, y).float().mean().data.item()

            loss = loss_var.data.item()
            cum_acc += acc
            cum_loss += loss
            
            if self.print_every and i % self.print_every == 0:
                print('[*] Batch %d, Loss: %.3f, Acc: %.3f' % (i, loss, acc))
        
        return cum_loss / n_batches, cum_acc / n_batches

    def save_checkpoint(self, filename='checkpoint.pth.tar'):
        state = {
            'model': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict()}
        torch.save(state, filename)

    def load_checkpoint(self, filename='checkpoint.pth.tar'):
        if os.path.isfile(filename):
            state = torch.load(filename)
            self.model.load_state_dict(state['model'])
            self.optimizer.load_state_dict(state['optimizer'])
        else:
            print('%s not found.' % filename)