Homec4science
Diffusion TIGraNet (master) Browse

R6062/utils.py
7553da4562f7master

Last Change
Raw File

utils.py
View Options

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Utilitary functions.
"""
import numpy as np
import matplotlib.pyplot as plt
import random
import glob
import logging
import torch
import torchvision
from torch.autograd import Variable
from torch.utils.data.sampler import SubsetRandomSampler
from configuration import *
random.seed(SEED)
logger = logging.getLogger(__name__)
def select(dataset, size, digits_to_keep, stratified_sampling=False):
"""Select randomly specific elements given by digits_to_keep."""
len_dataset = len(dataset)
indices = list(range(len_dataset))
random_select_indices = []
random.shuffle(indices)
if stratified_sampling:
num_classes = len(digits_to_keep)
classes = [[] for _ in range(num_classes)]
for i in indices:
if dataset[i][1] in digits_to_keep:
classe = dataset[i][1]
classes[classe].append(i)
for i in range(np.min([len(classes[0]), len(classes[1]), len(classes[2])])):
for j in range(num_classes):
if len(random_select_indices) < size:
random_select_indices.append(classes[j][i])
else:
break
else:
for i in indices:
if len(random_select_indices) < size and dataset[i][1] in digits_to_keep:
random_select_indices.append(i)
return random_select_indices
def train_valid_split(dataset, train_size):
"""Split the dataset into training and validaiton set."""
len_dataset = len(dataset)
indices = list(range(len_dataset))
train_indices = indices[:train_size]
valid_indices = indices[train_size:]
return SubsetRandomSampler(train_indices), SubsetRandomSampler(valid_indices)
def train_valid_test_split(dataset, train_size, valid_size):
"""Split the dataset into training, validation and testing set."""
indices = list(range(len(dataset)))
random.shuffle(indices)
train_indices, valid_indices, test_indices = indices[:train_size], indices[train_size:train_size+valid_size], indices[train_size+valid_size:]
return SubsetRandomSampler(train_indices), SubsetRandomSampler(valid_indices), SubsetRandomSampler(test_indices)
def imshow_data_loader(data_loader, eth80_class_names=[]):
"""Show image provided by the data loader."""
# get a batch of data
inputs, classes = next(iter(data_loader))
out = torchvision.utils.make_grid(tensor=inputs)
# get the corresponding values
if eth80_class_names:
title = [eth80_class_names[x] for x in classes]
mean = ETH80_MEAN
std = ETH80_STD
else:
title = [x for x in classes]
mean = MNIST_MEAN
std = MNIST_STD
# build the original image
out = out.numpy().transpose((1, 2, 0))
out = std * out + mean
out = np.clip(out, 0, 1)
# display it
plt.imshow(out)
plt.title(title)
plt.show()
def show_spectrum(tensor, num_filters):
"""Show the spectrum of the spectral layer. """
return NotImplemented
def snapshot(saved_model_dir, run_time, run_name, is_best, epoch, err_epoch, model_state_dict, optim_state_dict):
"""Save the model state."""
complete_name = '{}{}_{}_{}_{:.2f}'.format(saved_model_dir, run_time, run_name, epoch, err_epoch)
states = {
'model': model_state_dict,
'optimizer': optim_state_dict
}
# Save the model
with open(complete_name + '.pt', 'wb') as f:
torch.save(states, f)
def load_pretrained_model(saved_model_dir, run_name, model):
"""Load the specified model."""
states = glob.glob(saved_model_dir + run_name)[0]
if torch.cuda.is_available():
checkpoint = torch.load(states)
else:
checkpoint = torch.load(states, map_location=lambda storage, loc: storage)
model.load_state_dict(checkpoint['model'])
model.optimizer.load_state_dict(checkpoint['optimizer'])
logging.info('Loaded {} model.'.format(run_name))
return model
def init_mask(num_nodes, batch_size):
"""Initialize the nodes of interest by including all the nodes of the graph."""
mask = Variable(torch.ones(batch_size, num_nodes, 1)).to(DEVICE)
return mask
def count_class_freq(loader, num_classes):
"""Return the frequency for each class from the loader."""
t = np.zeros(num_classes)
for _, target in loader:
for c in target:
t[c] +=1
return t
def get_dim(data):
"""Get the dimension of the input image."""
dim = len(data[0])
return dim
c4science · Help