utils.py

from __future__ import absolute_import, division, print_function, unicode_literals

import collections
import math

import numpy as np
import os
import re
import scipy
import torch
import torch.utils
import torchvision
from io import BytesIO
from operator import itemgetter
from torch import nn
from torch.autograd import Variable
from torch.nn import init
from torch.utils.data import Dataset
from torchvision.transforms import transforms

from autoaugment import Cutout_, AutoAugment
from glico_model.model import weights_init
from vgg_arch import vgg19_bn
from wide_resnet import WideResNet
from wideresnet_2 import WideResNet2

PACKAGE_PARENT = '..'
SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
os.sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
from glico_model import vgg_metric

try:
	import tensorflow as tf
except:
	pass

NAGParams = collections.namedtuple('NAGParams',
                                   'nz force_l2 is_pixel z_init is_classifier disc_net loss data_name noise_proj shot')
NAGParams.__new__.__defaults__ = (None, None, None, None)
GANParams = collections.namedtuple('GANParams', 'ndf weight_d')
GANParams.__new__.__defaults__ = (None, None, None)
OptParams = collections.namedtuple('OptParams', 'lr factor ' +
                                   'batch_size epochs ' +
                                   'decay_epochs decay_rate')
OptParams.__new__.__defaults__ = (None, None, None, None, None, None)
ImageParams = collections.namedtuple('ImageParams', 'sz nc n mu sd')
ImageParams.__new__.__defaults__ = (None, None, None)


def distance_metric(sz, nc, force_l2=False):
	# return vgg_metric._VGGFixedDistance()
	if force_l2:
		return maybe_cuda(nn.L1Loss(), is_block=False)
	if sz == 16:
		return vgg_metric._VGGDistance(2)
	elif sz == 32:
		return vgg_metric._VGGDistance(3)
	elif sz == 64:
		return vgg_metric._VGGDistance(4)
	elif sz > 64:
		return vgg_metric._VGGMSDistance()


def _conv_layer(n_input, n_output, k, stride=1, padding=1, bias=True):
	"3x3 convolution with padding"
	seq = nn.Sequential(nn.Conv2d(n_input, n_output, kernel_size=k, stride=stride, padding=padding, bias=bias),
	                    nn.BatchNorm2d(n_output), nn.LeakyReLU(True),
	                    nn.MaxPool2d(kernel_size=2, stride=2, return_indices=False))
	# if Config.model.use_dropout:  # Add dropout module
	#     list_seq = list(seq.modules())[1:]
	#     list_seq.append(nn.Dropout(Config.model.dropout))
	#     seq = nn.Sequential(*list_seq)
	return seq


import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt


def plot_kernels(tensor, rn, epoch, num_cols=6):
	# Normalise
	maxVal = tensor.max()
	minVal = abs(tensor.min())
	maxVal = max(maxVal, minVal)
	tensor = tensor / maxVal
	tensor = tensor / 2
	tensor = tensor + 0.5
	num_rows = 1
	fig = plt.figure(figsize=(num_cols, num_rows))
	i = 0
	for t in tensor:
		ax1 = fig.add_subplot(num_rows, num_cols, i + 1)
		pilTrans = transforms.ToPILImage()
		pilImg = pilTrans(t)
		ax1.imshow(pilImg, interpolation='none')
		ax1.axis('off')
		ax1.set_xticklabels([])
		ax1.set_yticklabels([])
		i += 1

	plt.subplots_adjust(wspace=0.1, hspace=0.1)
	f'runs/ims_{rn}/conv1_{epoch}.png'


class cnn(nn.Module):
	"""
	A model for Mini-ImageNet classification.
	"""

	def __init__(self, num_classes, dim=64, im_size=84):
		super(cnn, self).__init__()
		self.arch = "vgg_small"
		self.n_filters = dim
		# The height and width of downsampled image
		ds_size = im_size // 2 ** 4

		self.layer1 = _conv_layer(3, self.n_filters, 3)
		# self.bn =  nn.BatchNorm2d(3)
		# self.pixel_norm = PixelNorm()
		# self.drop = nn.Dropout2d(drop_rate)
		self.lrelu = nn.LeakyReLU(0.1)
		self.layer2 = _conv_layer(self.n_filters, self.n_filters, 3)
		self.layer3 = _conv_layer(self.n_filters, self.n_filters, 3)
		self.layer4 = _conv_layer(self.n_filters, self.n_filters, 3)
		self.out = nn.Sequential(nn.Linear(self.n_filters * ds_size ** 2, num_classes))
		self.softmax = nn.Softmax()
		# self.out = nn.Linear(self.n_filters * ds_size ** 2, Config.model.n_classes)

		# Initialize layers
		self.weights_init(self.layer1)
		self.weights_init(self.layer2)
		self.weights_init(self.layer3)
		self.weights_init(self.layer4)

	def weights_init(self, module):
		for m in module.modules():
			if isinstance(m, nn.Conv2d):
				init.xavier_uniform_(m.weight, gain=np.sqrt(2))
				init.constant_(m.bias, 0)
			elif isinstance(m, nn.BatchNorm2d):
				m.weight.data.fill_(1)
				m.bias.data.zero_()

	def forward(self, x):
		# x = self.bn(x)
		# x=self.pixel_norm(x)
		x = self.layer1(x)
		x = self.lrelu(x)
		# x = self.drop(x)

		x = self.layer2(x)
		x = self.lrelu(x)
		# x = self.bn(x)
		# x = self.drop

		x = self.layer3(x)
		x = self.lrelu(x)
		# x = self.bn(x)

		x = self.layer4(x)
		x = self.lrelu(x)
		# x = self.drop(x)

		x = x.view(x.size(0), -1)
		x = self.out(x)
		return self.softmax(x)


def classifier(num_classes, dim=64):
	return cnn(num_classes=num_classes, dim=dim)


def sample_gaussian(x, m):
	x = x.data.numpy()
	mu = x.mean(0).squeeze()
	cov2 = np.cov(x, rowvar=0)
	z = np.random.multivariate_normal(mu, cov2, size=m)
	z_t = torch.from_numpy(z).float()
	radius = z_t.norm(2, 1).unsqueeze(1).expand_as(z_t)
	z_t = z_t / radius
	return Variable(maybe_cuda(z_t))


def maybe_cuda(tensor):
	return tensor.cuda() if torch.cuda.is_available() else tensor


class IndexToImageDataset(Dataset):
	"""Wrap a dataset to map indices to images
	In other words, instead of producing (X, y) it produces (idx, X). The label
	y is not relevant for our task.
	"""

	def __init__(self, base_dataset, transform=None, offset_idx=0, offset_label=0):
		'''

		:param base_dataset:
		:param transform:
		:param path2idx: verify the index is consistence to the path
		'''
		# assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors)
		self.base = base_dataset
		self.transform = transform
		self.offset_idx = offset_idx
		self.offset_label = offset_label

	def __len__(self):
		return len(self.base)

	def __getitem__(self, idx):
		# img, _ = self.base[idx]
		input = self.base[idx]
		label = input[1]
		# path = input[2]
		# print(f"{img[0].shape}, {type(img[0])}")
		if self.transform:
			img = self.transform(input[0])
		else:
			img = input[0]
		return (idx + self.offset_idx, img, label + self.offset_label)


class AverageMeter(object):
	"""Computes and stores the average and current value"""

	def __init__(self, name, fmt=':f'):
		self.name = name
		self.fmt = fmt
		self.reset()

	def reset(self):
		self.val = 0
		self.avg = 0
		self.sum = 0
		self.count = 0
		self.std = 0

	# self.all = []

	def update(self, val, n=1):
		# self.all.append(val)
		self.val = val
		self.sum += val * n
		self.count += n
		self.avg = self.sum / self.count

	# self.std = sqrt(mean(abs(val - np.mean(self.all)) ** 2))
	# self.sderr = 1.96 * self.std / np.sqrt(self.count)

	def __str__(self):
		fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
		return fmtstr.format(**self.__dict__)


class ProgressMeter(object):
	def __init__(self, num_batches, *meters, prefix=""):
		self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
		self.meters = meters
		self.prefix = prefix

	def print(self, batch):
		entries = [self.prefix + self.batch_fmtstr.format(batch)]
		entries += [str(meter) for meter in self.meters]
		print('\t'.join(entries))

	def _get_batch_fmtstr(self, num_batches):
		num_digits = len(str(num_batches // 1))
		fmt = '{:' + str(num_digits) + 'd}'
		return '[' + fmt + '/' + fmt.format(num_batches) + ']'


# import GPUtil as GPU
#
#
# GPUs = GPU.getGPUs()
#
#
# def printm(gpu):
# 	import psutil
# 	import humanize
# 	process = psutil.Process(os.getpid())
# 	print("Gen RAM Free: " + humanize.naturalsize(psutil.virtual_memory().available), " | Proc size: " + humanize.naturalsize(process.memory_info().rss))
# 	print("GPU RAM Free: {0:.0f}MB | Used: {1:.0f}MB | Util {2:3.0f}% | Total {3:.0f}MB".format(gpu.memoryFree, gpu.memoryUsed, gpu.memoryUtil * 100, gpu.memoryTotal))

def make_image_grid(x, ngrid):
	x = x.clone().cpu()
	if pow(ngrid, 2) < x.size(0):
		grid = make_grid(x[:ngrid * ngrid], nrow=ngrid, padding=0, normalize=True, scale_each=False)
	else:
		grid = torch.FloatTensor(ngrid * ngrid, x.size(1), x.size(2), x.size(3)).fill_(1)
		grid[:x.size(0)].copy_(x)
		grid = make_grid(grid, nrow=ngrid, padding=0, normalize=True, scale_each=False)
	return grid


def save_image_single(x, path, imsize=512):
	from PIL import Image
	grid = make_image_grid(x, 1)
	ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()
	im = Image.fromarray(ndarr)
	im = im.resize((imsize, imsize), Image.NEAREST)
	im.save(path)


def save_image_grid(x, path, imsize=512, ngrid=8):
	from PIL import Image
	grid = make_image_grid(x, ngrid)
	ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()
	im = Image.fromarray(ndarr)
	im = im.resize((imsize, imsize), Image.NEAREST)
	im.save(path)


irange = range


def make_grid(tensor, nrow=8, padding=2,
              normalize=False, range=None, scale_each=False, pad_value=0):
	"""Make a grid of images.
	Args:
		tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)
			or a list of images all of the same size.
		nrow (int, optional): Number of images displayed in each row of the grid.
			The Final grid size is (B / nrow, nrow). Default is 8.
		padding (int, optional): amount of padding. Default is 2.
		normalize (bool, optional): If True, shift the image to the range (0, 1),
			by subtracting the minimum and dividing by the maximum pixel value.
		range (tuple, optional): tuple (min, max) where min and max are numbers,
			then these numbers are used to normalize the image. By default, min and max
			are computed from the tensor.
		scale_each (bool, optional): If True, scale each image in the batch of
			images separately rather than the (min, max) over all images.
		pad_value (float, optional): Value for the padded pixels.
	Example:
		See this notebook `here <https://gist.github.com/anonymous/bf16430f7750c023141c562f3e9f2a91>`_
	"""
	if not (torch.is_tensor(tensor) or
	        (isinstance(tensor, list) and all(torch.is_tensor(t) for t in tensor))):
		raise TypeError('tensor or list of tensors expected, got {}'.format(type(tensor)))

	# if list of tensors, convert to a 4D mini-batch Tensor
	if isinstance(tensor, list):
		tensor = torch.stack(tensor, dim=0)

	if tensor.dim() == 2:  # single image H x W
		tensor = tensor.view(1, tensor.size(0), tensor.size(1))
	if tensor.dim() == 3:  # single image
		if tensor.size(0) == 1:  # if single-channel, convert to 3-channel
			tensor = torch.cat((tensor, tensor, tensor), 0)
		return tensor
	if tensor.dim() == 4 and tensor.size(1) == 1:  # single-channel images
		tensor = torch.cat((tensor, tensor, tensor), 1)

	if normalize is True:
		tensor = tensor.clone()  # avoid modifying tensor in-place
		if range is not None:
			assert isinstance(range, tuple), \
				"range has to be a tuple (min, max) if specified. min and max are numbers"

		def norm_ip(img, min, max):
			img.clamp_(min=min, max=max)
			img.add_(-min).div_(max - min)

		def norm_range(t, range):
			if range is not None:
				norm_ip(t, range[0], range[1])
			else:
				norm_ip(t, t.min(), t.max())

		if scale_each is True:
			for t in tensor:  # loop over mini-batch dimension
				norm_range(t, range)
		else:
			norm_range(tensor, range)

	# make the mini-batch of images into a grid
	nmaps = tensor.size(0)
	xmaps = min(nrow, nmaps)
	ymaps = int(math.ceil(float(nmaps) / xmaps))
	height, width = int(tensor.size(2) + padding), int(tensor.size(3) + padding)
	grid = tensor.new(3, height * ymaps + padding, width * xmaps + padding).fill_(pad_value)
	k = 0
	for y in irange(ymaps):
		for x in irange(xmaps):
			if k >= nmaps:
				break
			grid.narrow(1, y * height + padding, height - padding) \
				.narrow(2, x * width + padding, width - padding) \
				.copy_(tensor[k])
			k = k + 1
	return grid


def save_image(tensor, filename, nrow=8, padding=2,
               normalize=False, range=None, scale_each=False, pad_value=0):
	"""Save a given Tensor into an image file.
	Args:
		tensor (Tensor or list): Image to be saved. If given a mini-batch tensor,
			saves the tensor as a grid of images by calling ``make_grid``.
		**kwargs: Other arguments are documented in ``make_grid``.
	"""
	from PIL import Image
	tensor = tensor.cpu()
	grid = make_grid(tensor, nrow=nrow, padding=padding, pad_value=pad_value,
	                 normalize=normalize, range=range, scale_each=scale_each)
	ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).numpy()
	im = Image.fromarray(ndarr)
	im.save(filename)


def find_latest(find_path):
	sorted_path = get_sorted_path(find_path)
	if len(sorted_path) == 0:
		return None
	return sorted_path[-1]


def get_sorted_path(find_path):
	dir_path = os.path.dirname(find_path)
	base_name = os.path.basename(find_path)
	paths = []
	for root, dirs, files in os.walk(dir_path):
		for f_name in files:
			if f_name.startswith(base_name) and f_name.endswith(".pkl"):
				paths.append(os.path.join(root, f_name))
	return sorted(paths, key=lambda x: int(re.findall("\d+", os.path.basename(x))[0]))


def load_saved_model(path, model):
	latest_path = find_latest(path)
	if latest_path is None:
		print(f"No model has been found! in {path}")
		return 0, model
	if torch.cuda.is_available():
		checkpoint = torch.load(latest_path)
	else:
		checkpoint = torch.load(latest_path, map_location='cpu')
	epoch = checkpoint['epoch']
	state_dict = checkpoint['model']

	try:
		model.load_state_dict(checkpoint['model'])
	except Exception as e:
		print(e)
		# import traceback
		# print(traceback.format_exc())
		new_state_dict = collections.OrderedDict()
		for k, v in state_dict.items():
			name = k[7:]  # remove 'module.' of dataparallel
			new_state_dict[name] = v

		model.load_state_dict(new_state_dict)
	try:
		model.label2idx = checkpoint['label2idx']
		model.idx2label = checkpoint['idx2label']
	except:
		pass
	print(f"Load checkpoints...! {latest_path}")
	return epoch, model


def save_checkpoint(path_to_save, epoch, model, maps=None, max_to_keep=1):
	sorted_path = get_sorted_path(path_to_save)
	for i in range(len(sorted_path) - max_to_keep):
		os.remove(sorted_path[i])
	full_path = f"{path_to_save}_{epoch}.pkl"
	if maps is not None:
		torch.save({"epoch": epoch, 'model': model.state_dict(), "label2idx": maps[0], "idx2label": maps[1]}, full_path)
	else:
		torch.save({"epoch": epoch, 'model': model.state_dict()}, full_path)
	print(f"Save checkpoints...! {full_path}")


class TensorBoard:
	def __init__(self, log_dir):
		"""Create a summary writer logging to log_dir."""
		self.writer = tf.summary.FileWriter(log_dir)

	def scalar_summary(self, tag, value, step):
		"""Log a scalar variable."""
		summary = tf.Summary(value=[tf.Summary.Value(tag=tag, simple_value=value)])
		self.writer.add_summary(summary, step)

	def image_summary(self, tag, images, step):
		"""Log a list of images."""

		img_summaries = []
		for i, img in enumerate(images):
			# Write the image to a string
			s = BytesIO()
			scipy.misc.toimage(img).save(s, format="png")

			# Create an Image object
			img_sum = tf.Summary.Image(encoded_image_string=s.getvalue(), height=img.shape[0], width=img.shape[1])
			# Create a Summary value
			img_summaries.append(tf.Summary.Value(tag=f"{tag}/{i}", image=img_sum))

		# Create and write Summary
		summary = tf.Summary(value=img_summaries)
		self.writer.add_summary(summary, step)

	def histo_summary(self, tag, values, step, bins=1000):
		"""Log a histogram of the tensor of values."""

		# Create a histogram using numpy
		counts, bin_edges = np.histogram(values, bins=bins)

		# Fill the fields of the histogram proto
		hist = tf.HistogramProto()
		hist.min = float(np.min(values))
		hist.max = float(np.max(values))
		hist.num = int(np.prod(values.shape))
		hist.sum = float(np.sum(values))
		hist.sum_squares = float(np.sum(values ** 2))

		# Drop the start of the first bin
		bin_edges = bin_edges[1:]

		# Add bin edges and counts
		for edge in bin_edges:
			hist.bucket_limit.append(edge)
		for c in counts:
			hist.bucket.append(c)

		# Create and write Summary
		summary = tf.Summary(value=[tf.Summary.Value(tag=tag, histo=hist)])
		self.writer.add_summary(summary, step)
		self.writer.flush()

	def _add_summary(self, step, summary):
		for tag, value in summary.items():
			self.scalar_summary(tag, value, step)


def get_labels(dataset):
	max_label = 100
	if hasattr(dataset, 'classes'):
		return dataset.classes
	if hasattr(dataset, 'train_labels'):
		return dataset.train_labels
	if hasattr(dataset, 'labels'):
		return dataset.labels
	if hasattr(dataset, 'targets'):
		return dataset.targets
	if hasattr(dataset, 'test_labels'):
		return dataset.test_labels
	else:
		print('No labels found! ')
	return max_label


def _load_lowshot_cifar(split, data_dir, num_shot):
	"""Load mini-imagenet from numpy's npz file format."""
	split_tag = {'train': f"train_{num_shot}_shot", 'test': f"test_{num_shot}_shot"}[split]
	dataset_path = os.path.join(data_dir, 'small-shot-{}_v3.npz'.format(split_tag))
	data = np.load(dataset_path)
	fields = data['features'], data['targets']
	print(len(data['targets']))
	print(set(data['targets']))
	return fields


def get_cub_param():
	crop_size = 224  # 448
	target_size = 256  # 512
	mean = [0.485, 0.456, 0.406]
	std = [0.229, 0.224, 0.225]
	return {'std': std, 'mean': mean, 'rand_crop': crop_size, 'image_size': target_size}


def get_cifar_param():
	crop_size = 32
	target_size = 32
	mean = [0.507, 0.487, 0.441]
	std = [0.267, 0.256, 0.276]
	return {'std': std, 'mean': mean, 'rand_crop': crop_size, 'image_size': target_size}


def get_imagenet_param():
	crop_size = 84
	target_size = 96
	mean = [0.485, 0.456, 0.406]
	std = [0.229, 0.224, 0.225]
	return {'std': std, 'mean': mean, 'rand_crop': crop_size, 'image_size': target_size}


def get_train_stl_param():
	crop_size = 96
	target_size = 96
	mean = [0.44671062, 0.43980984, 0.40664645]
	std = [0.26034098, 0.25657727, 0.27126738]
	return {'std': std, 'mean': mean, 'rand_crop': crop_size, 'image_size': target_size}


def get_train_n_unlabled_stl_param():
	crop_size = 96
	target_size = 96
	mean = [0.44087802, 0.42790631, 0.38678794]
	std = [0.26826769, 0.26104504, 0.26866837]
	return {'std': std, 'mean': mean, 'rand_crop': crop_size, 'image_size': target_size}


def get_test_stl_param():
	crop_size = 96
	target_size = 96
	mean = [0.44723063, 0.43964247, 0.40495725]
	std = [0.2605645, 0.25666146, 0.26997382]
	return {'std': std, 'mean': mean, 'rand_crop': crop_size, 'image_size': target_size}


class normalize_np(object):
	def __init__(self, mean, std):
		self.mean, self.std = [np.array(a, np.float32) for a in (mean, std)]

	def __call__(self, x):
		x -= self.mean * 255
		x *= 1.0 / (255 * self.std)
		return x

	def __repr__(self):
		return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)


def transpose(x, source='NHWC', target='NCHW'):
	return x.transpose([source.index(d) for d in target])


def pad(x, border=4):
	return np.pad(x, [(0, 0), (border, border), (border, border)], mode='reflect')


class RandomPadandCrop(object):
	"""Crop randomly the image.

	Args:
		output_size (tuple or int): Desired output size. If int, square crop
			is made.
	"""

	def __init__(self, output_size):
		assert isinstance(output_size, (int, tuple))
		if isinstance(output_size, int):
			self.output_size = (output_size, output_size)
		else:
			assert len(output_size) == 2
			self.output_size = output_size

	def __call__(self, x):
		x = pad(x, 4)

		h, w = x.shape[1:]
		new_h, new_w = self.output_size

		top = np.random.randint(0, h - new_h)
		left = np.random.randint(0, w - new_w)

		x = x[:, top: top + new_h, left: left + new_w]

		return x


class RandomFlip(object):
	"""Flip randomly the image.
	"""

	def __call__(self, x):
		if np.random.rand() < 0.5:
			x = x[:, :, ::-1]

		return x.copy()


class GaussianNoise(object):
	"""Add gaussian noise to the image.
	"""

	def __call__(self, x):
		c, h, w = x.shape
		x += np.random.randn(c, h, w) * 0.15
		return x


class ToTensor(object):
	"""Transform the image to tensor.
	"""

	def __call__(self, x):
		x = torch.from_numpy(x)
		return x


def get_loader_with_idx(dataset, batch_size, image_size, rand_crop, mean=None, std=None, num_workers=6,
                        augment=False, shuffle=True,
                        offset_idx=0,
                        offset_label=0, sampler=None, eval=False, autoaugment=False, drop_last=False, cutout=False,
                        random_erase=False):
	'''
	Note not to use normalize in NagTrainer
	'''
	if std is None:
		std = [0.267, 0.256, 0.276]
	if mean is None:
		mean = [0.507, 0.487, 0.441]
	if sampler is not None:
		shuffle = False
	normalize = transforms.Normalize(mean=mean, std=std)  # CIFAR100
	transform_list = []

	is_cifar = image_size == 32
	is_stl = image_size == 96

	if isinstance(augment, bool):
		if augment:
			print(f"augment:{augment}")
			if is_cifar:
				transform_list.append(transforms.ToPILImage())  # Comment this linr for full data
				transform_list.append(transforms.RandomCrop(rand_crop, padding=4))
				transform_list.append(transforms.RandomHorizontalFlip(p=0.5))

			elif is_stl:
				transform_list.append(transforms.RandomCrop(image_size, padding=12))
				transform_list.append(transforms.RandomHorizontalFlip(p=0.5))
			else:  # CUB, Places365,imagenet
				transform_list.append(transforms.RandomResizedCrop(rand_crop, scale=(0.875, 1.)))
				transform_list.append(transforms.RandomHorizontalFlip(p=0.5))
			if autoaugment:
				transform_list.append(AutoAugment())
			if cutout:
				transform_list.append(Cutout_())
			transform_list.append(transforms.ToTensor())

		if augment:
			transform_list.append(normalize)
			if random_erase:
				transform_list.append(transforms.RandomErasing())
		elif eval:
			print(f"eval:{eval}")
			shuffle = False
			if not is_cifar:
				transform_list.append(transforms.Resize(image_size))
				transform_list.append(transforms.CenterCrop(rand_crop))
			transform_list.append(transforms.ToTensor())
			transform_list.append(normalize)
		else:  # not eval | augment is False:
			print(f"not eval | augment : TRUE")
			if not is_cifar:
				transform_list.append(transforms.Resize((image_size, image_size)))
			# transform_list.append(transforms.CenterCrop(rand_crop))
			transform_list.append(transforms.ToTensor())

		transform = transforms.Compose(transform_list)
	else:
		transform = None

	loader = torch.utils.data.DataLoader(
			IndexToImageDataset(dataset, transform=transform, offset_idx=offset_idx, offset_label=offset_label),
			batch_size=batch_size,
			shuffle=shuffle, num_workers=num_workers, pin_memory=False, sampler=sampler, drop_last=drop_last)

	print(f"=>Generated data loader, res={image_size}, workers={num_workers} transform={transform} sampler={sampler}")
	return loader


def one_hot(labels, n_classes):
	y = labels.unsqueeze(1).long()
	# One hot encoding buffer that you create out of the loop and just keep reusing
	y_onehot = torch.FloatTensor(len(y), n_classes).cuda()

	y_onehot.zero_()
	return y_onehot.scatter_(1, y, 1)


def validate_loader_consistency(netZ, idx):
	try:
		targets = itemgetter(*idx.numpy().tolist())(netZ.idx2label)
	except:
		print("Something is wrong with idx2labeldict - maybe netZ wan't loaded")
	return targets


# def get_cifar100_small(data_dir, shot):
# 	train_data_imgs, train_lables = _load_lowshot_cifar(data_dir=data_dir, split="train",
# 	                                                    num_shot=shot)
# 	train_data = CompatibleDataset(train_data_imgs, train_lables)
# 	test_data_fewshot_imgs, test_lables = _load_lowshot_cifar(data_dir=data_dir, split="test",
# 	                                                          num_shot=shot)
# 	transductive_train_data = CompatibleDataset(test_data_fewshot_imgs, test_lables)
#
# 	return train_data, transductive_train_data


def get_classifier(classes, d, pretrained):
	global cnn_
	feature_extracting = False

	def set_parameter_requires_grad(model, feature_extracting):
		if feature_extracting:
			for param in model.parameters():
				param.requires_grad = False

	if d == "vgg":
		cnn_ = vgg19_bn(num_classes=classes)
	elif d == "wideresnet":
		cnn_ = WideResNet(depth=28, num_classes=classes, widen_factor=10, dropRate=0.3)
	elif d == "wideresnet2":
		cnn_ = WideResNet2(num_classes=classes)
	# cnn_ = torch.hub.load('pytorch/vision:v0.4.2', 'wide_resnet101_2', pretrained=True)
	elif d == 'densenet':
		cnn_ = torchvision.models.densenet121(pretrained=pretrained)
		num_ftrs = cnn_.classifier.in_features
		cnn_.classifier = nn.Linear(num_ftrs, classes)
	elif d == "resnet":
		cnn_ = torchvision.models.resnet101(pretrained=pretrained)
		set_parameter_requires_grad(cnn_, feature_extracting)
		num_ftrs = cnn_.fc.in_features
		cnn_.fc = nn.Linear(num_ftrs, classes)
	elif d == "resnet50":
		cnn_ = torchvision.models.resnet50(pretrained=pretrained)
		set_parameter_requires_grad(cnn_, feature_extracting)
		num_ftrs = cnn_.fc.in_features
		cnn_.fc = nn.Linear(num_ftrs, classes)
	elif d == "resnet18":
		cnn_ = torchvision.models.resnet18(pretrained=pretrained)
		set_parameter_requires_grad(cnn_, feature_extracting)
		num_ftrs = cnn_.fc.in_features
		cnn_.fc = nn.Linear(num_ftrs, classes)
	elif d == "conv":
		cnn_ = cnn(num_classes=classes, im_size=32)  # Pixel space
		cnn_.apply(weights_init)
	return cnn_


class AugmentGaussian:
	def __init__(self, validation_stddev=25, train_stddev_rng_range=(0, 50)):
		self.validation_stddev = validation_stddev
		self.train_stddev_range = train_stddev_rng_range

	def add_gaussian(self, x):
		shape = x.size()
		(minval, maxval) = self.train_stddev_range
		train_stddev = np.random.uniform(minval / 255.0, maxval / 255.0, size=shape)
		noise = torch.randn_like(x).float().cuda()
		return x + noise * torch.tensor(train_stddev).float().cuda()

	def add_validation_noise(self, x):
		noise = torch.randn_like(x).float().cuda()
		return x + noise * torch.tensor(self.validation_stddev / 255.0).float().cuda()


class LabelSmoothingLoss(nn.Module):
	def __init__(self, classes, smoothing=0.0, dim=-1):
		super(LabelSmoothingLoss, self).__init__()
		self.confidence = 1.0 - smoothing
		self.smoothing = smoothing
		self.cls = classes
		self.dim = dim

	def forward(self, pred, target):
		pred = pred.log_softmax(dim=self.dim)
		with torch.no_grad():
			# true_dist = pred.data.clone()
			true_dist = torch.zeros_like(pred)
			true_dist.fill_(self.smoothing / (self.cls - 1))
			true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
		return torch.mean(torch.sum(-true_dist * pred, dim=self.dim))