"""Helpful functions for ATDA."""

import os
import random
from math import log

import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.autograd import Variable

from datasets import get_mnist, get_mnist_m, get_svhn, get_usps
from misc import config as cfg


def make_variable(tensor, volatile=False):
    """Convert Tensor to Variable."""
    if torch.cuda.is_available():
        tensor = tensor.cuda()
    return Variable(tensor, volatile=volatile)


def make_cuda(tensor):
    """Use CUDA if it's available."""
    if torch.cuda.is_available():
        tensor = tensor.cuda()
    return tensor


def denormalize(x, std, mean):
    """Invert normalization, and then convert array into image."""
    out = x * std + mean
    return out.clamp(0, 1)


def init_weights(layer):
    """Init weights for layers."""
    layer_name = layer.__class__.__name__
    if layer_name.find("Conv") != -1:
        layer.weight.data.normal_(0.0, 0.02)
    elif layer_name.find("BatchNorm") != -1:
        layer.weight.data.normal_(1.0, 0.02)
        layer.bias.data.fill_(0)


def init_random_seed(manual_seed):
    """Init random seed."""
    seed = None
    if manual_seed is None:
        seed = random.randint(1, 10000)
    else:
        seed = manual_seed
    print("use random seed: {}".format(seed))
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def enable_cudnn_benchmark():
    """Turn on the cudnn autotuner that selects efficient algorithms."""
    if torch.cuda.is_available():
        cudnn.benchmark = True


def init_model(net, restore):
    """Init models with cuda and weights."""
    # init weights of model
    net.apply(init_weights)

    # restore model weights
    restore_model(net, restore)

    # check if cuda is available
    if torch.cuda.is_available():
        net.cuda()

    return net


def save_model(net, filename):
    """Save trained model."""
    if not os.path.exists(cfg.model_root):
        os.makedirs(cfg.model_root)
    torch.save(net.state_dict(),
               os.path.join(cfg.model_root, filename))
    print("save pretrained model to: {}".format(os.path.join(cfg.model_root,
                                                             filename)))


def restore_model(net, restore):
    """Restore network from saved model."""
    if restore is not None and os.path.exists(restore):
        net.load_state_dict(torch.load(restore))
        net.restored = True
        print("Restore model from: {}".format(os.path.abspath(restore)))


def get_optimizer(net, name="Adam"):
    """Get optimizer by name."""
    if name == "Adam":
        return optim.Adam(net.parameters(),
                          lr=cfg.learning_rate,
                          betas=(cfg.beta1, cfg.beta2))


def get_data_loader(name, train=True, get_dataset=False):
    """Get data loader by name."""
    if name == "MNIST":
        return get_mnist(train, get_dataset)
    elif name == "MNIST-M":
        return get_mnist_m(train, get_dataset)
    elif name == "SVHN":
        return get_svhn(train, get_dataset)
    elif name == "USPS":
        return get_usps(train, get_dataset)


def make_data_loader(dataset, batch_size=cfg.batch_size,
                     shuffle=True, sampler=None):
    """Make dataloader from dataset."""
    data_loader = torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=shuffle if sampler is None else False,
        sampler=sampler)
    return data_loader


def get_inf_iterator(data_loader):
    """Inf data iterator."""
    while True:
        for images, labels in data_loader:
            yield (images, labels)


def get_model_params(net, name):
    """Get parameters of models by name."""
    for n, p in net.named_parameters():
        if n == name:
            return p


def calc_similiar_penalty(F_1, F_2):
    """Calculate similiar penalty |W_1^T W_2|."""
    F_1_params = get_model_params(F_1, "classifier.8.weight")
    F_2_params = get_model_params(F_2, "classifier.8.weight")
    similiar_penalty = torch.sum(
        torch.abs(torch.mm(F_1_params.transpose(0, 1), F_2_params)))
    return similiar_penalty


# def get_whole_dataset(dataset):
#     """Get all images and labels of dataset."""
#     data_loader = torch.utils.data.DataLoader(dataset=dataset,
#                                               batch_size=len(dataset))
#     for images, labels in data_loader:
#         return images, labels


# def expand_single_channel(data):
#     """Expand single channel images into three channels."""
#     if data.dim() == 4 and data.size(1) == 1:
#         return torch.cat([data, data, data], 1)
#     else:
#         return data

# No need for this, just use torch.utils.data.ConcatDataset
# def concat_dataset(images_a, labels_a, images_b, labels_b):
#     """Concatenate images and labels of two datasets."""
#     # ensure the same size of images_a and images_b
#     # images_a = expand_single_channel(images_a)
#     # images_b = expand_single_channel(images_b)
#     # concatenate images and labels
#     images = torch.cat([images_a, images_b], 0)
#     labels = torch.cat([labels_a, labels_b], 0)
#
#     return images, labels


# Just use get_sampled_data_loader()
# def sample_candidatas(images, labels, candidates_num, shuffle=True):
#     """Sample images and labels from dataset."""
#     # get indices
#     indices = torch.arange(0, len(images))
#     if shuffle:
#         indices = torch.randperm(len(images))
#     # slice indices
#     candidates_num = min(len(images), candidates_num)
#     excerpt = indices.narrow(0, 0, candidates_num).long()
#     # select items by indices
#     images_sampled = images.index_select(0, excerpt)
#     labels_sampled = labels.index_select(0, excerpt)
#     return images_sampled, labels_sampled


def get_sampled_data_loader(dataset, candidates_num, shuffle=True):
    """Get data loader for sampled dataset."""
    # get indices
    indices = torch.arange(0, len(dataset))
    if shuffle:
        indices = torch.randperm(len(dataset))
    # slice indices
    candidates_num = min(len(dataset), candidates_num)
    excerpt = indices.narrow(0, 0, candidates_num).long()
    sampler = torch.utils.data.sampler.SubsetRandomSampler(excerpt)
    return make_data_loader(dataset, sampler=sampler, shuffle=False)

# Just use torch.utils.data.DataLoader for iteration
# def get_minibatch_iterator(images, labels, batchsize, shuffle=False):
#     """Get minibatch iterator with given images and labels."""
#     assert len(images) == len(labels), \
#         "Number of images and labels must be equal to make minibatches!"
#
#     if shuffle:
#         indices = torch.randperm(len(images))
#
#     for start_idx in range(0, len(images), batchsize):
#         end_idx = start_idx + batchsize
#         if end_idx > len(images):
#             end_idx = start_idx + (len(images) % batchsize)
#
#         if shuffle:
#             excerpt = indices.narrow(0, start_idx, end_idx)
#         else:
#             excerpt = torch.arange(start_idx, end_idx).long()
#
#         images_batch = images.index_select(0, excerpt)
#         labels_batch = labels.index_select(0, excerpt)
#
#         yield images_batch, labels_batch


# No need to convert dense labels into one-hot labels in this experiment
# def make_labels(labels_dense, num_classes):
#     """Convert dense labels into one-hot labels."""
#     labels_one_hot = torch.zeros((labels_dense.size(0), num_classes))
#     labels_one_hot.scatter_(1, labels_dense, 1)
#     return labels_one_hot


def guess_pseudo_labels(out_1, out_2, threshold=0.9):
    """Guess labels of target dataset by the two outputs."""
    # get prediction
    _, pred_idx_1 = torch.max(out_1, 1)
    _, pred_idx_2 = torch.max(out_2, 1)
    # find prediction who are the same in two outputs
    equal_idx = torch.nonzero(torch.eq(pred_idx_1, pred_idx_2)).squeeze()

    try:
        out_1 = out_1[equal_idx, :]
        out_2 = out_2[equal_idx, :]
    except IndexError:
        print("out_1: {}".format(out_1.size()))
        print("out_2: {}".format(out_2.size()))
        print("equal_idx: {}".format(equal_idx.size()))
        out_1 = out_1[equal_idx, :]
        out_2 = out_2[equal_idx, :]

    # filter indices by threshold
    # note that we use log(threshold) since the output is LogSoftmax
    pred_1, _ = torch.max(out_1, 1)
    pred_2, _ = torch.max(out_2, 1)
    max_pred, _ = torch.max(torch.stack([pred_1, pred_2], 1), 1)
    filtered_idx = torch.nonzero(max_pred > log(threshold)).squeeze()
    # get images, pseudo labels and true labels by indices
    _, pred_idx = torch.max(out_1[filtered_idx, :], 1)

    pseudo_labels = pred_idx
    excerpt = equal_idx[filtered_idx]

    return excerpt, pseudo_labels

"""Helpful functions for ATDA."""

import os
import random
from math import log

import torch
import torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.autograd import Variable

from datasets import get_mnist, get_mnist_m, get_svhn, get_usps
from misc import config as cfg


def make_variable(tensor, volatile=False):
    """Convert Tensor to Variable."""
    if torch.cuda.is_available():
        tensor = tensor.cuda()
    return Variable(tensor, volatile=volatile)


def make_cuda(tensor):
    """Use CUDA if it's available."""
    if torch.cuda.is_available():
        tensor = tensor.cuda()
    return tensor


def denormalize(x, std, mean):
    """Invert normalization, and then convert array into image."""
    out = x * std + mean
    return out.clamp(0, 1)


def init_weights(layer):
    """Init weights for layers."""
    layer_name = layer.__class__.__name__
    if layer_name.find("Conv") != -1:
        layer.weight.data.normal_(0.0, 0.02)
    elif layer_name.find("BatchNorm") != -1:
        layer.weight.data.normal_(1.0, 0.02)
        layer.bias.data.fill_(0)


def init_random_seed(manual_seed):
    """Init random seed."""
    seed = None
    if manual_seed is None:
        seed = random.randint(1, 10000)
    else:
        seed = manual_seed
    print("use random seed: {}".format(seed))
    random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def enable_cudnn_benchmark():
    """Turn on the cudnn autotuner that selects efficient algorithms."""
    if torch.cuda.is_available():
        cudnn.benchmark = True


def init_model(net, restore):
    """Init models with cuda and weights."""
    # init weights of model
    net.apply(init_weights)

    # restore model weights
    restore_model(net, restore)

    # check if cuda is available
    if torch.cuda.is_available():
        net.cuda()

    return net


def save_model(net, filename):
    """Save trained model."""
    if not os.path.exists(cfg.model_root):
        os.makedirs(cfg.model_root)
    torch.save(net.state_dict(),
               os.path.join(cfg.model_root, filename))
    print("save pretrained model to: {}".format(os.path.join(cfg.model_root,
                                                             filename)))


def restore_model(net, restore):
    """Restore network from saved model."""
    if restore is not None and os.path.exists(restore):
        net.load_state_dict(torch.load(restore))
        net.restored = True
        print("Restore model from: {}".format(os.path.abspath(restore)))


def get_optimizer(net, name="Adam"):
    """Get optimizer by name."""
    if name == "Adam":
        return optim.Adam(net.parameters(),
                          lr=cfg.learning_rate,
                          betas=(cfg.beta1, cfg.beta2))


def get_data_loader(name, train=True, get_dataset=False):
    """Get data loader by name."""
    if name == "MNIST":
        return get_mnist(train, get_dataset)
    elif name == "MNIST-M":
        return get_mnist_m(train, get_dataset)
    elif name == "SVHN":
        return get_svhn(train, get_dataset)
    elif name == "USPS":
        return get_usps(train, get_dataset)


def make_data_loader(dataset, batch_size=cfg.batch_size,
                     shuffle=True, sampler=None):
    """Make dataloader from dataset."""
    data_loader = torch.utils.data.DataLoader(
        dataset=dataset,
        batch_size=batch_size,
        shuffle=shuffle if sampler is None else False,
        sampler=sampler)
    return data_loader


def get_inf_iterator(data_loader):
    """Inf data iterator."""
    while True:
        for images, labels in data_loader:
            yield (images, labels)


def get_model_params(net, name):
    """Get parameters of models by name."""
    for n, p in net.named_parameters():
        if n == name:
            return p


def calc_similiar_penalty(F_1, F_2):
    """Calculate similiar penalty |W_1^T W_2|."""
    F_1_params = get_model_params(F_1, "classifier.8.weight")
    F_2_params = get_model_params(F_2, "classifier.8.weight")
    similiar_penalty = torch.sum(
        torch.abs(torch.mm(F_1_params.transpose(0, 1), F_2_params)))
    return similiar_penalty


# def get_whole_dataset(dataset):
#     """Get all images and labels of dataset."""
#     data_loader = torch.utils.data.DataLoader(dataset=dataset,
#                                               batch_size=len(dataset))
#     for images, labels in data_loader:
#         return images, labels


# def expand_single_channel(data):
#     """Expand single channel images into three channels."""
#     if data.dim() == 4 and data.size(1) == 1:
#         return torch.cat([data, data, data], 1)
#     else:
#         return data

# No need for this, just use torch.utils.data.ConcatDataset
# def concat_dataset(images_a, labels_a, images_b, labels_b):
#     """Concatenate images and labels of two datasets."""
#     # ensure the same size of images_a and images_b
#     # images_a = expand_single_channel(images_a)
#     # images_b = expand_single_channel(images_b)
#     # concatenate images and labels
#     images = torch.cat([images_a, images_b], 0)
#     labels = torch.cat([labels_a, labels_b], 0)
#
#     return images, labels


# Just use get_sampled_data_loader()
# def sample_candidatas(images, labels, candidates_num, shuffle=True):
#     """Sample images and labels from dataset."""
#     # get indices
#     indices = torch.arange(0, len(images))
#     if shuffle:
#         indices = torch.randperm(len(images))
#     # slice indices
#     candidates_num = min(len(images), candidates_num)
#     excerpt = indices.narrow(0, 0, candidates_num).long()
#     # select items by indices
#     images_sampled = images.index_select(0, excerpt)
#     labels_sampled = labels.index_select(0, excerpt)
#     return images_sampled, labels_sampled


def get_sampled_data_loader(dataset, candidates_num, shuffle=True):
    """Get data loader for sampled dataset."""
    # get indices
    indices = torch.arange(0, len(dataset))
    if shuffle:
        indices = torch.randperm(len(dataset))
    # slice indices
    candidates_num = min(len(dataset), candidates_num)
    excerpt = indices.narrow(0, 0, candidates_num).long()
    sampler = torch.utils.data.sampler.SubsetRandomSampler(excerpt)
    return make_data_loader(dataset, sampler=sampler, shuffle=False)

# Just use torch.utils.data.DataLoader for iteration
# def get_minibatch_iterator(images, labels, batchsize, shuffle=False):
#     """Get minibatch iterator with given images and labels."""
#     assert len(images) == len(labels), \
#         "Number of images and labels must be equal to make minibatches!"
#
#     if shuffle:
#         indices = torch.randperm(len(images))
#
#     for start_idx in range(0, len(images), batchsize):
#         end_idx = start_idx + batchsize
#         if end_idx > len(images):
#             end_idx = start_idx + (len(images) % batchsize)
#
#         if shuffle:
#             excerpt = indices.narrow(0, start_idx, end_idx)
#         else:
#             excerpt = torch.arange(start_idx, end_idx).long()
#
#         images_batch = images.index_select(0, excerpt)
#         labels_batch = labels.index_select(0, excerpt)
#
#         yield images_batch, labels_batch


# No need to convert dense labels into one-hot labels in this experiment
# def make_labels(labels_dense, num_classes):
#     """Convert dense labels into one-hot labels."""
#     labels_one_hot = torch.zeros((labels_dense.size(0), num_classes))
#     labels_one_hot.scatter_(1, labels_dense, 1)
#     return labels_one_hot


def guess_pseudo_labels(out_1, out_2, threshold=0.9):
    """Guess labels of target dataset by the two outputs."""
    # get prediction
    _, pred_idx_1 = torch.max(out_1, 1)
    _, pred_idx_2 = torch.max(out_2, 1)
    # find prediction who are the same in two outputs
    equal_idx = torch.nonzero(torch.eq(pred_idx_1, pred_idx_2)).squeeze()

    try:
        out_1 = out_1[equal_idx, :]
        out_2 = out_2[equal_idx, :]
    except IndexError:
        print("out_1: {}".format(out_1.size()))
        print("out_2: {}".format(out_2.size()))
        print("equal_idx: {}".format(equal_idx.size()))
        out_1 = out_1[equal_idx, :]
        out_2 = out_2[equal_idx, :]

    # filter indices by threshold
    # note that we use log(threshold) since the output is LogSoftmax
    pred_1, _ = torch.max(out_1, 1)
    pred_2, _ = torch.max(out_2, 1)
    max_pred, _ = torch.max(torch.stack([pred_1, pred_2], 1), 1)
    filtered_idx = torch.nonzero(max_pred > log(threshold)).squeeze()
    # get images, pseudo labels and true labels by indices
    _, pred_idx = torch.max(out_1[filtered_idx, :], 1)

    pseudo_labels = pred_idx
    excerpt = equal_idx[filtered_idx]

    return excerpt, pseudo_labels

Helpful functions for ATDA. Convert Tensor to Variable. Use CUDA if it's available. Invert normalization, and then convert array into image. Init weights for layers. Init random seed. Turn on the cudnn autotuner that selects efficient algorithms. Init models with cuda and weights. # init weights of model # restore model weights # check if cuda is available Save trained model. Restore network from saved model. Get optimizer by name. Get data loader by name. Make dataloader from dataset. Inf data iterator. Get parameters of models by name. Calculate similiar penalty |W_1^T W_2|. # def get_whole_dataset(dataset): #     """Get all images and labels of dataset.""" #     data_loader = torch.utils.data.DataLoader(dataset=dataset, #                                               batch_size=len(dataset)) #     for images, labels in data_loader: #         return images, labels # def expand_single_channel(data): #     """Expand single channel images into three channels.""" #     if data.dim() == 4 and data.size(1) == 1: #         return torch.cat([data, data, data], 1) #     else: #         return data # No need for this, just use torch.utils.data.ConcatDataset # def concat_dataset(images_a, labels_a, images_b, labels_b): #     """Concatenate images and labels of two datasets.""" #     # ensure the same size of images_a and images_b #     # images_a = expand_single_channel(images_a) #     # images_b = expand_single_channel(images_b) #     # concatenate images and labels #     images = torch.cat([images_a, images_b], 0) #     labels = torch.cat([labels_a, labels_b], 0) # #     return images, labels # Just use get_sampled_data_loader() # def sample_candidatas(images, labels, candidates_num, shuffle=True): #     """Sample images and labels from dataset.""" #     # get indices #     indices = torch.arange(0, len(images)) #     if shuffle: #         indices = torch.randperm(len(images)) #     # slice indices #     candidates_num = min(len(images), candidates_num) #     excerpt = indices.narrow(0, 0, candidates_num).long() #     # select items by indices #     images_sampled = images.index_select(0, excerpt) #     labels_sampled = labels.index_select(0, excerpt) #     return images_sampled, labels_sampled Get data loader for sampled dataset. # get indices # slice indices # Just use torch.utils.data.DataLoader for iteration # def get_minibatch_iterator(images, labels, batchsize, shuffle=False): #     """Get minibatch iterator with given images and labels.""" #     assert len(images) == len(labels), \ #         "Number of images and labels must be equal to make minibatches!" # #     if shuffle: #         indices = torch.randperm(len(images)) # #     for start_idx in range(0, len(images), batchsize): #         end_idx = start_idx + batchsize #         if end_idx > len(images): #             end_idx = start_idx + (len(images) % batchsize) # #         if shuffle: #             excerpt = indices.narrow(0, start_idx, end_idx) #         else: #             excerpt = torch.arange(start_idx, end_idx).long() # #         images_batch = images.index_select(0, excerpt) #         labels_batch = labels.index_select(0, excerpt) # #         yield images_batch, labels_batch # No need to convert dense labels into one-hot labels in this experiment # def make_labels(labels_dense, num_classes): #     """Convert dense labels into one-hot labels.""" #     labels_one_hot = torch.zeros((labels_dense.size(0), num_classes)) #     labels_one_hot.scatter_(1, labels_dense, 1) #     return labels_one_hot Guess labels of target dataset by the two outputs. # get prediction # find prediction who are the same in two outputs # filter indices by threshold # note that we use log(threshold) since the output is LogSoftmax # get images, pseudo labels and true labels by indices