Move to xautodl

xautodl/utils/__init__.py

@@ -0,0 +1,13 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019.01 #
+#####################################################
+# This directory contains some ad-hoc functions, classes, etc.
+# It will be re-formulated in the future.
+#####################################################
+from .evaluation_utils import obtain_accuracy
+from .gpu_manager import GPUManager
+from .flop_benchmark import get_model_infos, count_parameters, count_parameters_in_MB
+from .affine_utils import normalize_points, denormalize_points
+from .affine_utils import identity2affine, solve2theta, affine2image
+from .hash_utils import get_md5_file
+from .str_utils import split_str2indexes

xautodl/utils/affine_utils.py

@@ -0,0 +1,159 @@
+# functions for affine transformation
+import math
+import torch
+import numpy as np
+import torch.nn.functional as F
+
+
+def identity2affine(full=False):
+    if not full:
+        parameters = torch.zeros((2, 3))
+        parameters[0, 0] = parameters[1, 1] = 1
+    else:
+        parameters = torch.zeros((3, 3))
+        parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1
+    return parameters
+
+
+def normalize_L(x, L):
+    return -1.0 + 2.0 * x / (L - 1)
+
+
+def denormalize_L(x, L):
+    return (x + 1.0) / 2.0 * (L - 1)
+
+
+def crop2affine(crop_box, W, H):
+    assert len(crop_box) == 4, "Invalid crop-box : {:}".format(crop_box)
+    parameters = torch.zeros(3, 3)
+    x1, y1 = normalize_L(crop_box[0], W), normalize_L(crop_box[1], H)
+    x2, y2 = normalize_L(crop_box[2], W), normalize_L(crop_box[3], H)
+    parameters[0, 0] = (x2 - x1) / 2
+    parameters[0, 2] = (x2 + x1) / 2
+
+    parameters[1, 1] = (y2 - y1) / 2
+    parameters[1, 2] = (y2 + y1) / 2
+    parameters[2, 2] = 1
+    return parameters
+
+
+def scale2affine(scalex, scaley):
+    parameters = torch.zeros(3, 3)
+    parameters[0, 0] = scalex
+    parameters[1, 1] = scaley
+    parameters[2, 2] = 1
+    return parameters
+
+
+def offset2affine(offx, offy):
+    parameters = torch.zeros(3, 3)
+    parameters[0, 0] = parameters[1, 1] = parameters[2, 2] = 1
+    parameters[0, 2] = offx
+    parameters[1, 2] = offy
+    return parameters
+
+
+def horizontalmirror2affine():
+    parameters = torch.zeros(3, 3)
+    parameters[0, 0] = -1
+    parameters[1, 1] = parameters[2, 2] = 1
+    return parameters
+
+
+# clockwise rotate image = counterclockwise rotate the rectangle
+# degree is between [0, 360]
+def rotate2affine(degree):
+    assert degree >= 0 and degree <= 360, "Invalid degree : {:}".format(degree)
+    degree = degree / 180 * math.pi
+    parameters = torch.zeros(3, 3)
+    parameters[0, 0] = math.cos(-degree)
+    parameters[0, 1] = -math.sin(-degree)
+    parameters[1, 0] = math.sin(-degree)
+    parameters[1, 1] = math.cos(-degree)
+    parameters[2, 2] = 1
+    return parameters
+
+
+# shape is a tuple [H, W]
+def normalize_points(shape, points):
+    assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(
+        shape
+    ) == 2, "invalid shape : {:}".format(shape)
+    assert isinstance(points, torch.Tensor) and (
+        points.shape[0] == 2
+    ), "points are wrong : {:}".format(points.shape)
+    (H, W), points = shape, points.clone()
+    points[0, :] = normalize_L(points[0, :], W)
+    points[1, :] = normalize_L(points[1, :], H)
+    return points
+
+
+# shape is a tuple [H, W]
+def normalize_points_batch(shape, points):
+    assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(
+        shape
+    ) == 2, "invalid shape : {:}".format(shape)
+    assert isinstance(points, torch.Tensor) and (
+        points.size(-1) == 2
+    ), "points are wrong : {:}".format(points.shape)
+    (H, W), points = shape, points.clone()
+    x = normalize_L(points[..., 0], W)
+    y = normalize_L(points[..., 1], H)
+    return torch.stack((x, y), dim=-1)
+
+
+# shape is a tuple [H, W]
+def denormalize_points(shape, points):
+    assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(
+        shape
+    ) == 2, "invalid shape : {:}".format(shape)
+    assert isinstance(points, torch.Tensor) and (
+        points.shape[0] == 2
+    ), "points are wrong : {:}".format(points.shape)
+    (H, W), points = shape, points.clone()
+    points[0, :] = denormalize_L(points[0, :], W)
+    points[1, :] = denormalize_L(points[1, :], H)
+    return points
+
+
+# shape is a tuple [H, W]
+def denormalize_points_batch(shape, points):
+    assert (isinstance(shape, tuple) or isinstance(shape, list)) and len(
+        shape
+    ) == 2, "invalid shape : {:}".format(shape)
+    assert isinstance(points, torch.Tensor) and (
+        points.shape[-1] == 2
+    ), "points are wrong : {:}".format(points.shape)
+    (H, W), points = shape, points.clone()
+    x = denormalize_L(points[..., 0], W)
+    y = denormalize_L(points[..., 1], H)
+    return torch.stack((x, y), dim=-1)
+
+
+# make target * theta = source
+def solve2theta(source, target):
+    source, target = source.clone(), target.clone()
+    oks = source[2, :] == 1
+    assert torch.sum(oks).item() >= 3, "valid points : {:} is short".format(oks)
+    if target.size(0) == 2:
+        target = torch.cat((target, oks.unsqueeze(0).float()), dim=0)
+    source, target = source[:, oks], target[:, oks]
+    source, target = source.transpose(1, 0), target.transpose(1, 0)
+    assert source.size(1) == target.size(1) == 3
+    # X, residual, rank, s = np.linalg.lstsq(target.numpy(), source.numpy())
+    # theta = torch.Tensor(X.T[:2, :])
+    X_, qr = torch.gels(source, target)
+    theta = X_[:3, :2].transpose(1, 0)
+    return theta
+
+
+# shape = [H,W]
+def affine2image(image, theta, shape):
+    C, H, W = image.size()
+    theta = theta[:2, :].unsqueeze(0)
+    grid_size = torch.Size([1, C, shape[0], shape[1]])
+    grid = F.affine_grid(theta, grid_size)
+    affI = F.grid_sample(
+        image.unsqueeze(0), grid, mode="bilinear", padding_mode="border"
+    )
+    return affI.squeeze(0)

xautodl/utils/evaluation_utils.py

@@ -0,0 +1,17 @@
+import torch
+
+
+def obtain_accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res

xautodl/utils/flop_benchmark.py

@@ -0,0 +1,218 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+def count_parameters_in_MB(model):
+    return count_parameters(model, "mb")
+
+
+def count_parameters(model_or_parameters, unit="mb"):
+    if isinstance(model_or_parameters, nn.Module):
+        counts = sum(np.prod(v.size()) for v in model_or_parameters.parameters())
+    elif isinstance(models_or_parameters, nn.Parameter):
+        counts = models_or_parameters.numel()
+    elif isinstance(models_or_parameters, (list, tuple)):
+        counts = sum(count_parameters(x, None) for x in models_or_parameters)
+    else:
+        counts = sum(np.prod(v.size()) for v in model_or_parameters)
+    if unit.lower() == "kb" or unit.lower() == "k":
+        counts /= 2 ** 10  # changed from 1e3 to 2^10
+    elif unit.lower() == "mb" or unit.lower() == "m":
+        counts /= 2 ** 20  # changed from 1e6 to 2^20
+    elif unit.lower() == "gb" or unit.lower() == "g":
+        counts /= 2 ** 30  # changed from 1e9 to 2^30
+    elif unit is not None:
+        raise ValueError("Unknow unit: {:}".format(unit))
+    return counts
+
+
+def get_model_infos(model, shape):
+    # model = copy.deepcopy( model )
+
+    model = add_flops_counting_methods(model)
+    # model = model.cuda()
+    model.eval()
+
+    # cache_inputs = torch.zeros(*shape).cuda()
+    # cache_inputs = torch.zeros(*shape)
+    cache_inputs = torch.rand(*shape)
+    if next(model.parameters()).is_cuda:
+        cache_inputs = cache_inputs.cuda()
+    # print_log('In the calculating function : cache input size : {:}'.format(cache_inputs.size()), log)
+    with torch.no_grad():
+        _____ = model(cache_inputs)
+    FLOPs = compute_average_flops_cost(model) / 1e6
+    Param = count_parameters_in_MB(model)
+
+    if hasattr(model, "auxiliary_param"):
+        aux_params = count_parameters_in_MB(model.auxiliary_param())
+        print("The auxiliary params of this model is : {:}".format(aux_params))
+        print(
+            "We remove the auxiliary params from the total params ({:}) when counting".format(
+                Param
+            )
+        )
+        Param = Param - aux_params
+
+    # print_log('FLOPs : {:} MB'.format(FLOPs), log)
+    torch.cuda.empty_cache()
+    model.apply(remove_hook_function)
+    return FLOPs, Param
+
+
+# ---- Public functions
+def add_flops_counting_methods(model):
+    model.__batch_counter__ = 0
+    add_batch_counter_hook_function(model)
+    model.apply(add_flops_counter_variable_or_reset)
+    model.apply(add_flops_counter_hook_function)
+    return model
+
+
+def compute_average_flops_cost(model):
+    """
+    A method that will be available after add_flops_counting_methods() is called on a desired net object.
+    Returns current mean flops consumption per image.
+    """
+    batches_count = model.__batch_counter__
+    flops_sum = 0
+    # or isinstance(module, torch.nn.AvgPool2d) or isinstance(module, torch.nn.MaxPool2d) \
+    for module in model.modules():
+        if (
+            isinstance(module, torch.nn.Conv2d)
+            or isinstance(module, torch.nn.Linear)
+            or isinstance(module, torch.nn.Conv1d)
+            or hasattr(module, "calculate_flop_self")
+        ):
+            flops_sum += module.__flops__
+    return flops_sum / batches_count
+
+
+# ---- Internal functions
+def pool_flops_counter_hook(pool_module, inputs, output):
+    batch_size = inputs[0].size(0)
+    kernel_size = pool_module.kernel_size
+    out_C, output_height, output_width = output.shape[1:]
+    assert out_C == inputs[0].size(1), "{:} vs. {:}".format(out_C, inputs[0].size())
+
+    overall_flops = (
+        batch_size * out_C * output_height * output_width * kernel_size * kernel_size
+    )
+    pool_module.__flops__ += overall_flops
+
+
+def self_calculate_flops_counter_hook(self_module, inputs, output):
+    overall_flops = self_module.calculate_flop_self(inputs[0].shape, output.shape)
+    self_module.__flops__ += overall_flops
+
+
+def fc_flops_counter_hook(fc_module, inputs, output):
+    batch_size = inputs[0].size(0)
+    xin, xout = fc_module.in_features, fc_module.out_features
+    assert xin == inputs[0].size(1) and xout == output.size(1), "IO=({:}, {:})".format(
+        xin, xout
+    )
+    overall_flops = batch_size * xin * xout
+    if fc_module.bias is not None:
+        overall_flops += batch_size * xout
+    fc_module.__flops__ += overall_flops
+
+
+def conv1d_flops_counter_hook(conv_module, inputs, outputs):
+    batch_size = inputs[0].size(0)
+    outL = outputs.shape[-1]
+    [kernel] = conv_module.kernel_size
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+    conv_per_position_flops = kernel * in_channels * out_channels / groups
+
+    active_elements_count = batch_size * outL
+    overall_flops = conv_per_position_flops * active_elements_count
+
+    if conv_module.bias is not None:
+        overall_flops += out_channels * active_elements_count
+    conv_module.__flops__ += overall_flops
+
+
+def conv2d_flops_counter_hook(conv_module, inputs, output):
+    batch_size = inputs[0].size(0)
+    output_height, output_width = output.shape[2:]
+
+    kernel_height, kernel_width = conv_module.kernel_size
+    in_channels = conv_module.in_channels
+    out_channels = conv_module.out_channels
+    groups = conv_module.groups
+    conv_per_position_flops = (
+        kernel_height * kernel_width * in_channels * out_channels / groups
+    )
+
+    active_elements_count = batch_size * output_height * output_width
+    overall_flops = conv_per_position_flops * active_elements_count
+
+    if conv_module.bias is not None:
+        overall_flops += out_channels * active_elements_count
+    conv_module.__flops__ += overall_flops
+
+
+def batch_counter_hook(module, inputs, output):
+    # Can have multiple inputs, getting the first one
+    inputs = inputs[0]
+    batch_size = inputs.shape[0]
+    module.__batch_counter__ += batch_size
+
+
+def add_batch_counter_hook_function(module):
+    if not hasattr(module, "__batch_counter_handle__"):
+        handle = module.register_forward_hook(batch_counter_hook)
+        module.__batch_counter_handle__ = handle
+
+
+def add_flops_counter_variable_or_reset(module):
+    if (
+        isinstance(module, torch.nn.Conv2d)
+        or isinstance(module, torch.nn.Linear)
+        or isinstance(module, torch.nn.Conv1d)
+        or isinstance(module, torch.nn.AvgPool2d)
+        or isinstance(module, torch.nn.MaxPool2d)
+        or hasattr(module, "calculate_flop_self")
+    ):
+        module.__flops__ = 0
+
+
+def add_flops_counter_hook_function(module):
+    if isinstance(module, torch.nn.Conv2d):
+        if not hasattr(module, "__flops_handle__"):
+            handle = module.register_forward_hook(conv2d_flops_counter_hook)
+            module.__flops_handle__ = handle
+    elif isinstance(module, torch.nn.Conv1d):
+        if not hasattr(module, "__flops_handle__"):
+            handle = module.register_forward_hook(conv1d_flops_counter_hook)
+            module.__flops_handle__ = handle
+    elif isinstance(module, torch.nn.Linear):
+        if not hasattr(module, "__flops_handle__"):
+            handle = module.register_forward_hook(fc_flops_counter_hook)
+            module.__flops_handle__ = handle
+    elif isinstance(module, torch.nn.AvgPool2d) or isinstance(
+        module, torch.nn.MaxPool2d
+    ):
+        if not hasattr(module, "__flops_handle__"):
+            handle = module.register_forward_hook(pool_flops_counter_hook)
+            module.__flops_handle__ = handle
+    elif hasattr(module, "calculate_flop_self"):  # self-defined module
+        if not hasattr(module, "__flops_handle__"):
+            handle = module.register_forward_hook(self_calculate_flops_counter_hook)
+            module.__flops_handle__ = handle
+
+
+def remove_hook_function(module):
+    hookers = ["__batch_counter_handle__", "__flops_handle__"]
+    for hooker in hookers:
+        if hasattr(module, hooker):
+            handle = getattr(module, hooker)
+            handle.remove()
+    keys = ["__flops__", "__batch_counter__", "__flops__"] + hookers
+    for ckey in keys:
+        if hasattr(module, ckey):
+            delattr(module, ckey)

xautodl/utils/gpu_manager.py

@@ -0,0 +1,86 @@
+import os
+
+
+class GPUManager:
+    queries = (
+        "index",
+        "gpu_name",
+        "memory.free",
+        "memory.used",
+        "memory.total",
+        "power.draw",
+        "power.limit",
+    )
+
+    def __init__(self):
+        all_gpus = self.query_gpu(False)
+
+    def get_info(self, ctype):
+        cmd = "nvidia-smi --query-gpu={} --format=csv,noheader".format(ctype)
+        lines = os.popen(cmd).readlines()
+        lines = [line.strip("\n") for line in lines]
+        return lines
+
+    def query_gpu(self, show=True):
+        num_gpus = len(self.get_info("index"))
+        all_gpus = [{} for i in range(num_gpus)]
+        for query in self.queries:
+            infos = self.get_info(query)
+            for idx, info in enumerate(infos):
+                all_gpus[idx][query] = info
+
+        if "CUDA_VISIBLE_DEVICES" in os.environ:
+            CUDA_VISIBLE_DEVICES = os.environ["CUDA_VISIBLE_DEVICES"].split(",")
+            selected_gpus = []
+            for idx, CUDA_VISIBLE_DEVICE in enumerate(CUDA_VISIBLE_DEVICES):
+                find = False
+                for gpu in all_gpus:
+                    if gpu["index"] == CUDA_VISIBLE_DEVICE:
+                        assert not find, "Duplicate cuda device index : {}".format(
+                            CUDA_VISIBLE_DEVICE
+                        )
+                        find = True
+                        selected_gpus.append(gpu.copy())
+                        selected_gpus[-1]["index"] = "{}".format(idx)
+                assert find, "Does not find the device : {}".format(CUDA_VISIBLE_DEVICE)
+            all_gpus = selected_gpus
+
+        if show:
+            allstrings = ""
+            for gpu in all_gpus:
+                string = "| "
+                for query in self.queries:
+                    if query.find("memory") == 0:
+                        xinfo = "{:>9}".format(gpu[query])
+                    else:
+                        xinfo = gpu[query]
+                    string = string + query + " : " + xinfo + " | "
+                allstrings = allstrings + string + "\n"
+            return allstrings
+        else:
+            return all_gpus
+
+    def select_by_memory(self, numbers=1):
+        all_gpus = self.query_gpu(False)
+        assert numbers <= len(all_gpus), "Require {} gpus more than you have".format(
+            numbers
+        )
+        alls = []
+        for idx, gpu in enumerate(all_gpus):
+            free_memory = gpu["memory.free"]
+            free_memory = free_memory.split(" ")[0]
+            free_memory = int(free_memory)
+            index = gpu["index"]
+            alls.append((free_memory, index))
+        alls.sort(reverse=True)
+        alls = [int(alls[i][1]) for i in range(numbers)]
+        return sorted(alls)
+
+
+"""
+if __name__ == '__main__':
+  manager = GPUManager()
+  manager.query_gpu(True)
+  indexes = manager.select_by_memory(3)
+  print (indexes)
+"""

xautodl/utils/hash_utils.py

@@ -0,0 +1,17 @@
+import os
+import hashlib
+
+
+def get_md5_file(file_path, post_truncated=5):
+    md5_hash = hashlib.md5()
+    if os.path.exists(file_path):
+        xfile = open(file_path, "rb")
+        content = xfile.read()
+        md5_hash.update(content)
+        digest = md5_hash.hexdigest()
+    else:
+        raise ValueError("[get_md5_file] {:} does not exist".format(file_path))
+    if post_truncated is None:
+        return digest
+    else:
+        return digest[-post_truncated:]

xautodl/utils/nas_utils.py

@@ -0,0 +1,76 @@
+# This file is for experimental usage
+import torch, random
+import numpy as np
+from copy import deepcopy
+import torch.nn as nn
+
+# modules in AutoDL
+from models import CellStructure
+from log_utils import time_string
+
+
+def evaluate_one_shot(model, xloader, api, cal_mode, seed=111):
+    print(
+        "This is an old version of codes to use NAS-Bench-API, and should be modified to align with the new version. Please contact me for more details if you use this function."
+    )
+    weights = deepcopy(model.state_dict())
+    model.train(cal_mode)
+    with torch.no_grad():
+        logits = nn.functional.log_softmax(model.arch_parameters, dim=-1)
+        archs = CellStructure.gen_all(model.op_names, model.max_nodes, False)
+        probs, accuracies, gt_accs_10_valid, gt_accs_10_test = [], [], [], []
+        loader_iter = iter(xloader)
+        random.seed(seed)
+        random.shuffle(archs)
+        for idx, arch in enumerate(archs):
+            arch_index = api.query_index_by_arch(arch)
+            metrics = api.get_more_info(arch_index, "cifar10-valid", None, False, False)
+            gt_accs_10_valid.append(metrics["valid-accuracy"])
+            metrics = api.get_more_info(arch_index, "cifar10", None, False, False)
+            gt_accs_10_test.append(metrics["test-accuracy"])
+            select_logits = []
+            for i, node_info in enumerate(arch.nodes):
+                for op, xin in node_info:
+                    node_str = "{:}<-{:}".format(i + 1, xin)
+                    op_index = model.op_names.index(op)
+                    select_logits.append(logits[model.edge2index[node_str], op_index])
+            cur_prob = sum(select_logits).item()
+            probs.append(cur_prob)
+        cor_prob_valid = np.corrcoef(probs, gt_accs_10_valid)[0, 1]
+        cor_prob_test = np.corrcoef(probs, gt_accs_10_test)[0, 1]
+        print(
+            "{:} correlation for probabilities : {:.6f} on CIFAR-10 validation and {:.6f} on CIFAR-10 test".format(
+                time_string(), cor_prob_valid, cor_prob_test
+            )
+        )
+
+        for idx, arch in enumerate(archs):
+            model.set_cal_mode("dynamic", arch)
+            try:
+                inputs, targets = next(loader_iter)
+            except:
+                loader_iter = iter(xloader)
+                inputs, targets = next(loader_iter)
+            _, logits = model(inputs.cuda())
+            _, preds = torch.max(logits, dim=-1)
+            correct = (preds == targets.cuda()).float()
+            accuracies.append(correct.mean().item())
+            if idx != 0 and (idx % 500 == 0 or idx + 1 == len(archs)):
+                cor_accs_valid = np.corrcoef(accuracies, gt_accs_10_valid[: idx + 1])[
+                    0, 1
+                ]
+                cor_accs_test = np.corrcoef(accuracies, gt_accs_10_test[: idx + 1])[
+                    0, 1
+                ]
+                print(
+                    "{:} {:05d}/{:05d} mode={:5s}, correlation : accs={:.5f} for CIFAR-10 valid, {:.5f} for CIFAR-10 test.".format(
+                        time_string(),
+                        idx,
+                        len(archs),
+                        "Train" if cal_mode else "Eval",
+                        cor_accs_valid,
+                        cor_accs_test,
+                    )
+                )
+    model.load_state_dict(weights)
+    return archs, probs, accuracies

xautodl/utils/qlib_utils.py

@@ -0,0 +1,122 @@
+import os
+import numpy as np
+from typing import List, Text
+from collections import defaultdict, OrderedDict
+
+
+class QResult:
+    """A class to maintain the results of a qlib experiment."""
+
+    def __init__(self, name):
+        self._result = defaultdict(list)
+        self._name = name
+        self._recorder_paths = []
+        self._date2ICs = []
+
+    def append(self, key, value):
+        self._result[key].append(value)
+
+    def append_path(self, xpath):
+        self._recorder_paths.append(xpath)
+
+    def append_date2ICs(self, date2IC):
+        if self._date2ICs:  # not empty
+            keys = sorted(list(date2IC.keys()))
+            pre_keys = sorted(list(self._date2ICs[0].keys()))
+            assert len(keys) == len(pre_keys)
+            for i, (x, y) in enumerate(zip(keys, pre_keys)):
+                assert x == y, "[{:}] {:} vs {:}".format(i, x, y)
+        self._date2ICs.append(date2IC)
+
+    def find_all_dates(self):
+        dates = self._date2ICs[-1].keys()
+        return sorted(list(dates))
+
+    def get_IC_by_date(self, date, scale=1.0):
+        values = []
+        for date2IC in self._date2ICs:
+            values.append(date2IC[date] * scale)
+        return float(np.mean(values)), float(np.std(values))
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def paths(self):
+        return self._recorder_paths
+
+    @property
+    def result(self):
+        return self._result
+
+    @property
+    def keys(self):
+        return list(self._result.keys())
+
+    def __len__(self):
+        return len(self._result)
+
+    def __repr__(self):
+        return "{name}({xname}, {num} metrics)".format(
+            name=self.__class__.__name__, xname=self.name, num=len(self.result)
+        )
+
+    def __getitem__(self, key):
+        if key not in self._result:
+            raise ValueError(
+                "Invalid key {:}, please use one of {:}".format(key, self.keys)
+            )
+        values = self._result[key]
+        return float(np.mean(values))
+
+    def update(self, metrics, filter_keys=None):
+        for key, value in metrics.items():
+            if filter_keys is not None and key in filter_keys:
+                key = filter_keys[key]
+            elif filter_keys is not None:
+                continue
+            self.append(key, value)
+
+    @staticmethod
+    def full_str(xstr, space):
+        xformat = "{:" + str(space) + "s}"
+        return xformat.format(str(xstr))
+
+    @staticmethod
+    def merge_dict(dict_list):
+        new_dict = dict()
+        for xkey in dict_list[0].keys():
+            values = [x for xdict in dict_list for x in xdict[xkey]]
+            new_dict[xkey] = values
+        return new_dict
+
+    def info(
+        self,
+        keys: List[Text],
+        separate: Text = "& ",
+        space: int = 20,
+        verbose: bool = True,
+    ):
+        avaliable_keys = []
+        for key in keys:
+            if key not in self.result:
+                print("There are invalid key [{:}].".format(key))
+            else:
+                avaliable_keys.append(key)
+        head_str = separate.join([self.full_str(x, space) for x in avaliable_keys])
+        values = []
+        for key in avaliable_keys:
+            if "IR" in key:
+                current_values = [x * 100 for x in self._result[key]]
+            else:
+                current_values = self._result[key]
+            mean = np.mean(current_values)
+            std = np.std(current_values)
+            # values.append("{:.4f} $\pm$ {:.4f}".format(mean, std))
+            values.append("{:.2f} $\pm$ {:.2f}".format(mean, std))
+        value_str = separate.join([self.full_str(x, space) for x in values])
+        if verbose:
+            print(head_str)
+            print(value_str)
+        return head_str, value_str

xautodl/utils/str_utils.py

@@ -0,0 +1,21 @@
+def split_str2indexes(string: str, max_check: int, length_limit=5):
+    if not isinstance(string, str):
+        raise ValueError("Invalid scheme for {:}".format(string))
+    srangestr = "".join(string.split())
+    indexes = set()
+    for srange in srangestr.split(","):
+        srange = srange.split("-")
+        if len(srange) != 2:
+            raise ValueError("invalid srange : {:}".format(srange))
+        if length_limit is not None:
+            assert (
+                len(srange[0]) == len(srange[1]) == length_limit
+            ), "invalid srange : {:}".format(srange)
+        srange = (int(srange[0]), int(srange[1]))
+        if not (0 <= srange[0] <= srange[1] < max_check):
+            raise ValueError(
+                "{:} vs {:} vs {:}".format(srange[0], srange[1], max_check)
+            )
+        for i in range(srange[0], srange[1] + 1):
+            indexes.add(i)
+    return indexes

xautodl/utils/temp_sync.py

@@ -0,0 +1,63 @@
+# To be deleted.
+import copy
+import torch
+
+from xlayers.super_core import SuperSequential, SuperMLPv1
+from xlayers.super_core import SuperSimpleNorm
+from xlayers.super_core import SuperLinear
+
+
+def optimize_fn(xs, ys, device="cpu", max_iter=2000, max_lr=0.1):
+    xs = torch.FloatTensor(xs).view(-1, 1).to(device)
+    ys = torch.FloatTensor(ys).view(-1, 1).to(device)
+
+    model = SuperSequential(
+        SuperSimpleNorm(xs.mean().item(), xs.std().item()),
+        SuperLinear(1, 200),
+        torch.nn.LeakyReLU(),
+        SuperLinear(200, 100),
+        torch.nn.LeakyReLU(),
+        SuperLinear(100, 1),
+    ).to(device)
+    model.train()
+    optimizer = torch.optim.Adam(
+        model.parameters(), lr=max_lr, amsgrad=True
+    )
+    loss_func = torch.nn.MSELoss()
+    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+        optimizer,
+        milestones=[
+            int(max_iter * 0.25),
+            int(max_iter * 0.5),
+            int(max_iter * 0.75),
+        ],
+        gamma=0.3,
+    )
+
+    best_loss, best_param = None, None
+    for _iter in range(max_iter):
+        preds = model(xs)
+
+        optimizer.zero_grad()
+        loss = loss_func(preds, ys)
+        loss.backward()
+        optimizer.step()
+        lr_scheduler.step()
+
+        if best_loss is None or best_loss > loss.item():
+            best_loss = loss.item()
+            best_param = copy.deepcopy(model.state_dict())
+        
+        # print('loss={:}, best-loss={:}'.format(loss.item(), best_loss))
+    model.load_state_dict(best_param)
+    return model, loss_func, best_loss
+
+
+def evaluate_fn(model, xs, ys, loss_fn, device="cpu"):
+    with torch.no_grad():
+        inputs = torch.FloatTensor(xs).view(-1, 1).to(device)
+        ys = torch.FloatTensor(ys).view(-1, 1).to(device)
+        preds = model(inputs)
+        loss = loss_fn(preds, ys)
+        preds = preds.view(-1).cpu().numpy()
+    return preds, loss.item()

xautodl/utils/weight_watcher.py

@@ -0,0 +1,400 @@
+#####################################################
+# Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2020.03 #
+#####################################################
+# Reformulate the codes in https://github.com/CalculatedContent/WeightWatcher
+#####################################################
+import numpy as np
+from typing import List
+import torch.nn as nn
+from collections import OrderedDict
+from sklearn.decomposition import TruncatedSVD
+
+
+def available_module_types():
+    return (nn.Conv2d, nn.Linear)
+
+
+def get_conv2D_Wmats(tensor: np.ndarray) -> List[np.ndarray]:
+    """
+    Extract W slices from a 4 index conv2D tensor of shape: (N,M,i,j) or (M,N,i,j).
+    Return ij (N x M) matrices
+    """
+    mats = []
+    N, M, imax, jmax = tensor.shape
+    assert (
+        N + M >= imax + jmax
+    ), "invalid tensor shape detected: {}x{} (NxM), {}x{} (i,j)".format(
+        N, M, imax, jmax
+    )
+    for i in range(imax):
+        for j in range(jmax):
+            w = tensor[:, :, i, j]
+            if N < M:
+                w = w.T
+            mats.append(w)
+    return mats
+
+
+def glorot_norm_check(W, N, M, rf_size, lower=0.5, upper=1.5):
+    """Check if this layer needs Glorot Normalization Fix"""
+
+    kappa = np.sqrt(2 / ((N + M) * rf_size))
+    norm = np.linalg.norm(W)
+
+    check1 = norm / np.sqrt(N * M)
+    check2 = norm / (kappa * np.sqrt(N * M))
+
+    if (rf_size > 1) and (check2 > lower) and (check2 < upper):
+        return check2, True
+    elif (check1 > lower) & (check1 < upper):
+        return check1, True
+    else:
+        if rf_size > 1:
+            return check2, False
+        else:
+            return check1, False
+
+
+def glorot_norm_fix(w, n, m, rf_size):
+    """Apply Glorot Normalization Fix."""
+    kappa = np.sqrt(2 / ((n + m) * rf_size))
+    w = w / kappa
+    return w
+
+
+def analyze_weights(
+    weights,
+    min_size,
+    max_size,
+    alphas,
+    lognorms,
+    spectralnorms,
+    softranks,
+    normalize,
+    glorot_fix,
+):
+    results = OrderedDict()
+    count = len(weights)
+    if count == 0:
+        return results
+
+    for i, weight in enumerate(weights):
+        M, N = np.min(weight.shape), np.max(weight.shape)
+        Q = N / M
+        results[i] = cur_res = OrderedDict(N=N, M=M, Q=Q)
+        check, checkTF = glorot_norm_check(weight, N, M, count)
+        cur_res["check"] = check
+        cur_res["checkTF"] = checkTF
+        # assume receptive field size is count
+        if glorot_fix:
+            weight = glorot_norm_fix(weight, N, M, count)
+        else:
+            # probably never needed since we always fix for glorot
+            weight = weight * np.sqrt(count / 2.0)
+
+        if spectralnorms:  # spectralnorm is the max eigenvalues
+            svd = TruncatedSVD(n_components=1, n_iter=7, random_state=10)
+            svd.fit(weight)
+            sv = svd.singular_values_
+            sv_max = np.max(sv)
+            if normalize:
+                evals = sv * sv / N
+            else:
+                evals = sv * sv
+            lambda0 = evals[0]
+            cur_res["spectralnorm"] = lambda0
+            cur_res["logspectralnorm"] = np.log10(lambda0)
+        else:
+            lambda0 = None
+
+        if M < min_size:
+            summary = "Weight matrix {}/{} ({},{}): Skipping: too small (<{})".format(
+                i + 1, count, M, N, min_size
+            )
+            cur_res["summary"] = summary
+            continue
+        elif max_size > 0 and M > max_size:
+            summary = (
+                "Weight matrix {}/{} ({},{}): Skipping: too big (testing) (>{})".format(
+                    i + 1, count, M, N, max_size
+                )
+            )
+            cur_res["summary"] = summary
+            continue
+        else:
+            summary = []
+        if alphas:
+            import powerlaw
+
+            svd = TruncatedSVD(n_components=M - 1, n_iter=7, random_state=10)
+            svd.fit(weight.astype(float))
+            sv = svd.singular_values_
+            if normalize:
+                evals = sv * sv / N
+            else:
+                evals = sv * sv
+
+            lambda_max = np.max(evals)
+            fit = powerlaw.Fit(evals, xmax=lambda_max, verbose=False)
+            alpha = fit.alpha
+            cur_res["alpha"] = alpha
+            D = fit.D
+            cur_res["D"] = D
+            cur_res["lambda_min"] = np.min(evals)
+            cur_res["lambda_max"] = lambda_max
+            alpha_weighted = alpha * np.log10(lambda_max)
+            cur_res["alpha_weighted"] = alpha_weighted
+            tolerance = lambda_max * M * np.finfo(np.max(sv)).eps
+            cur_res["rank_loss"] = np.count_nonzero(sv > tolerance, axis=-1)
+
+            logpnorm = np.log10(np.sum([ev ** alpha for ev in evals]))
+            cur_res["logpnorm"] = logpnorm
+
+            summary.append(
+                "Weight matrix {}/{} ({},{}): Alpha: {}, Alpha Weighted: {}, D: {}, pNorm {}".format(
+                    i + 1, count, M, N, alpha, alpha_weighted, D, logpnorm
+                )
+            )
+
+        if lognorms:
+            norm = np.linalg.norm(weight)  # Frobenius Norm
+            cur_res["norm"] = norm
+            lognorm = np.log10(norm)
+            cur_res["lognorm"] = lognorm
+
+            X = np.dot(weight.T, weight)
+            if normalize:
+                X = X / N
+            normX = np.linalg.norm(X)  # Frobenius Norm
+            cur_res["normX"] = normX
+            lognormX = np.log10(normX)
+            cur_res["lognormX"] = lognormX
+
+            summary.append(
+                "Weight matrix {}/{} ({},{}): LogNorm: {} ; LogNormX: {}".format(
+                    i + 1, count, M, N, lognorm, lognormX
+                )
+            )
+
+            if softranks:
+                softrank = norm ** 2 / sv_max ** 2
+                softranklog = np.log10(softrank)
+                softranklogratio = lognorm / np.log10(sv_max)
+                cur_res["softrank"] = softrank
+                cur_res["softranklog"] = softranklog
+                cur_res["softranklogratio"] = softranklogratio
+                summary += (
+                    "{}. Softrank: {}. Softrank log: {}. Softrank log ratio: {}".format(
+                        summary, softrank, softranklog, softranklogratio
+                    )
+                )
+        cur_res["summary"] = "\n".join(summary)
+    return results
+
+
+def compute_details(results):
+    """
+    Return a pandas data frame.
+    """
+    final_summary = OrderedDict()
+
+    metrics = {
+        # key in "results" : pretty print name
+        "check": "Check",
+        "checkTF": "CheckTF",
+        "norm": "Norm",
+        "lognorm": "LogNorm",
+        "normX": "Norm X",
+        "lognormX": "LogNorm X",
+        "alpha": "Alpha",
+        "alpha_weighted": "Alpha Weighted",
+        "spectralnorm": "Spectral Norm",
+        "logspectralnorm": "Log Spectral Norm",
+        "softrank": "Softrank",
+        "softranklog": "Softrank Log",
+        "softranklogratio": "Softrank Log Ratio",
+        "sigma_mp": "Marchenko-Pastur (MP) fit sigma",
+        "numofSpikes": "Number of spikes per MP fit",
+        "ratio_numofSpikes": "aka, percent_mass, Number of spikes / total number of evals",
+        "softrank_mp": "Softrank for MP fit",
+        "logpnorm": "alpha pNorm",
+    }
+
+    metrics_stats = []
+    for metric in metrics:
+        metrics_stats.append("{}_min".format(metric))
+        metrics_stats.append("{}_max".format(metric))
+        metrics_stats.append("{}_avg".format(metric))
+
+        metrics_stats.append("{}_compound_min".format(metric))
+        metrics_stats.append("{}_compound_max".format(metric))
+        metrics_stats.append("{}_compound_avg".format(metric))
+
+    columns = (
+        [
+            "layer_id",
+            "layer_type",
+            "N",
+            "M",
+            "layer_count",
+            "slice",
+            "slice_count",
+            "level",
+            "comment",
+        ]
+        + [*metrics]
+        + metrics_stats
+    )
+
+    metrics_values = {}
+    metrics_values_compound = {}
+
+    for metric in metrics:
+        metrics_values[metric] = []
+        metrics_values_compound[metric] = []
+
+    layer_count = 0
+    for layer_id, result in results.items():
+        layer_count += 1
+
+        layer_type = np.NAN
+        if "layer_type" in result:
+            layer_type = str(result["layer_type"]).replace("LAYER_TYPE.", "")
+
+        compounds = {}  # temp var
+        for metric in metrics:
+            compounds[metric] = []
+
+        slice_count, Ntotal, Mtotal = 0, 0, 0
+        for slice_id, summary in result.items():
+            if not str(slice_id).isdigit():
+                continue
+            slice_count += 1
+
+            N = np.NAN
+            if "N" in summary:
+                N = summary["N"]
+                Ntotal += N
+
+            M = np.NAN
+            if "M" in summary:
+                M = summary["M"]
+                Mtotal += M
+
+            data = {
+                "layer_id": layer_id,
+                "layer_type": layer_type,
+                "N": N,
+                "M": M,
+                "slice": slice_id,
+                "level": "SLICE",
+                "comment": "Slice level",
+            }
+            for metric in metrics:
+                if metric in summary:
+                    value = summary[metric]
+                    if value is not None:
+                        metrics_values[metric].append(value)
+                        compounds[metric].append(value)
+                        data[metric] = value
+
+        data = {
+            "layer_id": layer_id,
+            "layer_type": layer_type,
+            "N": Ntotal,
+            "M": Mtotal,
+            "slice_count": slice_count,
+            "level": "LAYER",
+            "comment": "Layer level",
+        }
+        # Compute the compound value over the slices
+        for metric, value in compounds.items():
+            count = len(value)
+            if count == 0:
+                continue
+
+            compound = np.mean(value)
+            metrics_values_compound[metric].append(compound)
+            data[metric] = compound
+
+    data = {"layer_count": layer_count, "level": "NETWORK", "comment": "Network Level"}
+    for metric, metric_name in metrics.items():
+        if metric not in metrics_values or len(metrics_values[metric]) == 0:
+            continue
+
+        values = metrics_values[metric]
+        minimum = min(values)
+        maximum = max(values)
+        avg = np.mean(values)
+        final_summary[metric] = avg
+        # print("{}: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg))
+        data["{}_min".format(metric)] = minimum
+        data["{}_max".format(metric)] = maximum
+        data["{}_avg".format(metric)] = avg
+
+        values = metrics_values_compound[metric]
+        minimum = min(values)
+        maximum = max(values)
+        avg = np.mean(values)
+        final_summary["{}_compound".format(metric)] = avg
+        # print("{} compound: min: {}, max: {}, avg: {}".format(metric_name, minimum, maximum, avg))
+        data["{}_compound_min".format(metric)] = minimum
+        data["{}_compound_max".format(metric)] = maximum
+        data["{}_compound_avg".format(metric)] = avg
+
+    return final_summary
+
+
+def analyze(
+    model: nn.Module,
+    min_size=50,
+    max_size=0,
+    alphas: bool = False,
+    lognorms: bool = True,
+    spectralnorms: bool = False,
+    softranks: bool = False,
+    normalize: bool = False,
+    glorot_fix: bool = False,
+):
+    """
+    Analyze the weight matrices of a model.
+    :param model: A PyTorch model
+    :param min_size: The minimum weight matrix size to analyze.
+    :param max_size: The maximum weight matrix size to analyze (0 = no limit).
+    :param alphas: Compute the power laws (alpha) of the weight matrices.
+      Time consuming so disabled by default (use lognorm if you want speed)
+    :param lognorms: Compute the log norms of the weight matrices.
+    :param spectralnorms: Compute the spectral norm (max eigenvalue) of the weight matrices.
+    :param softranks: Compute the soft norm (i.e. StableRank) of the weight matrices.
+    :param normalize: Normalize or not.
+    :param glorot_fix:
+    :return: (a dict of all layers' results, a dict of the summarized info)
+    """
+    names, modules = [], []
+    for name, module in model.named_modules():
+        if isinstance(module, available_module_types()):
+            names.append(name)
+            modules.append(module)
+    # print('There are {:} layers to be analyzed in this model.'.format(len(modules)))
+    all_results = OrderedDict()
+    for index, module in enumerate(modules):
+        if isinstance(module, nn.Linear):
+            weights = [module.weight.cpu().detach().numpy()]
+        else:
+            weights = get_conv2D_Wmats(module.weight.cpu().detach().numpy())
+        results = analyze_weights(
+            weights,
+            min_size,
+            max_size,
+            alphas,
+            lognorms,
+            spectralnorms,
+            softranks,
+            normalize,
+            glorot_fix,
+        )
+        results["id"] = index
+        results["type"] = type(module)
+        all_results[index] = results
+    summary = compute_details(all_results)
+    return all_results, summary