wonkitty/apple_oh

1# Copyright (c) Facebook, Inc. and its affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7import errno8import functools9import hashlib10import inspect11import io12import os13import random14import socket15import tempfile16import warnings17import zlib18from contextlib import contextmanager19 20from diffq import UniformQuantizer, DiffQuantizer21import torch as th22import tqdm23from torch import distributed24from torch.nn import functional as F25 26 27def center_trim(tensor, reference):28    """29    Center trim `tensor` with respect to `reference`, along the last dimension.30    `reference` can also be a number, representing the length to trim to.31    If the size difference != 0 mod 2, the extra sample is removed on the right side.32    """33    if hasattr(reference, "size"):34        reference = reference.size(-1)35    delta = tensor.size(-1) - reference36    if delta < 0:37        raise ValueError("tensor must be larger than reference. " f"Delta is {delta}.")38    if delta:39        tensor = tensor[..., delta // 2:-(delta - delta // 2)]40    return tensor41 42 43def average_metric(metric, count=1.):44    """45    Average `metric` which should be a float across all hosts. `count` should be46    the weight for this particular host (i.e. number of examples).47    """48    metric = th.tensor([count, count * metric], dtype=th.float32, device='cuda')49    distributed.all_reduce(metric, op=distributed.ReduceOp.SUM)50    return metric[1].item() / metric[0].item()51 52 53def free_port(host='', low=20000, high=40000):54    """55    Return a port number that is most likely free.56    This could suffer from a race condition although57    it should be quite rare.58    """59    sock = socket.socket()60    while True:61        port = random.randint(low, high)62        try:63            sock.bind((host, port))64        except OSError as error:65            if error.errno == errno.EADDRINUSE:66                continue67            raise68        return port69 70 71def sizeof_fmt(num, suffix='B'):72    """73    Given `num` bytes, return human readable size.74    Taken from https://stackoverflow.com/a/109493375    """76    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:77        if abs(num) < 1024.0:78            return "%3.1f%s%s" % (num, unit, suffix)79        num /= 1024.080    return "%.1f%s%s" % (num, 'Yi', suffix)81 82 83def human_seconds(seconds, display='.2f'):84    """85    Given `seconds` seconds, return human readable duration.86    """87    value = seconds * 1e688    ratios = [1e3, 1e3, 60, 60, 24]89    names = ['us', 'ms', 's', 'min', 'hrs', 'days']90    last = names.pop(0)91    for name, ratio in zip(names, ratios):92        if value / ratio < 0.3:93            break94        value /= ratio95        last = name96    return f"{format(value, display)} {last}"97 98 99class TensorChunk:100    def __init__(self, tensor, offset=0, length=None):101        total_length = tensor.shape[-1]102        assert offset >= 0103        assert offset < total_length104 105        if length is None:106            length = total_length - offset107        else:108            length = min(total_length - offset, length)109 110        self.tensor = tensor111        self.offset = offset112        self.length = length113        self.device = tensor.device114 115    @property116    def shape(self):117        shape = list(self.tensor.shape)118        shape[-1] = self.length119        return shape120 121    def padded(self, target_length):122        delta = target_length - self.length123        total_length = self.tensor.shape[-1]124        assert delta >= 0125 126        start = self.offset - delta // 2127        end = start + target_length128 129        correct_start = max(0, start)130        correct_end = min(total_length, end)131 132        pad_left = correct_start - start133        pad_right = end - correct_end134 135        out = F.pad(self.tensor[..., correct_start:correct_end], (pad_left, pad_right))136        assert out.shape[-1] == target_length137        return out138 139 140def tensor_chunk(tensor_or_chunk):141    if isinstance(tensor_or_chunk, TensorChunk):142        return tensor_or_chunk143    else:144        assert isinstance(tensor_or_chunk, th.Tensor)145        return TensorChunk(tensor_or_chunk)146 147 148def apply_model(model, mix, shifts=None, split=False,149                overlap=0.25, transition_power=1., progress=False):150    """151    Apply model to a given mixture.152 153    Args:154        shifts (int): if > 0, will shift in time `mix` by a random amount between 0 and 0.5 sec155            and apply the oppositve shift to the output. This is repeated `shifts` time and156            all predictions are averaged. This effectively makes the model time equivariant157            and improves SDR by up to 0.2 points.158        split (bool): if True, the input will be broken down in 8 seconds extracts159            and predictions will be performed individually on each and concatenated.160            Useful for model with large memory footprint like Tasnet.161        progress (bool): if True, show a progress bar (requires split=True)162    """163    assert transition_power >= 1, "transition_power < 1 leads to weird behavior."164    device = mix.device165    channels, length = mix.shape166    if split:167        out = th.zeros(len(model.sources), channels, length, device=device)168        sum_weight = th.zeros(length, device=device)169        segment = model.segment_length170        stride = int((1 - overlap) * segment)171        offsets = range(0, length, stride)172        scale = stride / model.samplerate173        if progress:174            offsets = tqdm.tqdm(offsets, unit_scale=scale, ncols=120, unit='seconds')175        # We start from a triangle shaped weight, with maximal weight in the middle176        # of the segment. Then we normalize and take to the power `transition_power`.177        # Large values of transition power will lead to sharper transitions.178        weight = th.cat([th.arange(1, segment // 2 + 1),179                         th.arange(segment - segment // 2, 0, -1)]).to(device)180        assert len(weight) == segment181        # If the overlap < 50%, this will translate to linear transition when182        # transition_power is 1.183        weight = (weight / weight.max())**transition_power184        for offset in offsets:185            chunk = TensorChunk(mix, offset, segment)186            chunk_out = apply_model(model, chunk, shifts=shifts)187            chunk_length = chunk_out.shape[-1]188            out[..., offset:offset + segment] += weight[:chunk_length] * chunk_out189            sum_weight[offset:offset + segment] += weight[:chunk_length]190            offset += segment191        assert sum_weight.min() > 0192        out /= sum_weight193        return out194    elif shifts:195        max_shift = int(0.5 * model.samplerate)196        mix = tensor_chunk(mix)197        padded_mix = mix.padded(length + 2 * max_shift)198        out = 0199        for _ in range(shifts):200            offset = random.randint(0, max_shift)201            shifted = TensorChunk(padded_mix, offset, length + max_shift - offset)202            shifted_out = apply_model(model, shifted)203            out += shifted_out[..., max_shift - offset:]204        out /= shifts205        return out206    else:207        valid_length = model.valid_length(length)208        mix = tensor_chunk(mix)209        padded_mix = mix.padded(valid_length)210        with th.no_grad():211            out = model(padded_mix.unsqueeze(0))[0]212        return center_trim(out, length)213 214 215@contextmanager216def temp_filenames(count, delete=True):217    names = []218    try:219        for _ in range(count):220            names.append(tempfile.NamedTemporaryFile(delete=False).name)221        yield names222    finally:223        if delete:224            for name in names:225                os.unlink(name)226 227 228def get_quantizer(model, args, optimizer=None):229    quantizer = None230    if args.diffq:231        quantizer = DiffQuantizer(232            model, min_size=args.q_min_size, group_size=8)233        if optimizer is not None:234            quantizer.setup_optimizer(optimizer)235    elif args.qat:236        quantizer = UniformQuantizer(237                model, bits=args.qat, min_size=args.q_min_size)238    return quantizer239 240 241def load_model(path, strict=False):242    with warnings.catch_warnings():243        warnings.simplefilter("ignore")244        load_from = path245        package = th.load(load_from, 'cpu')246 247    klass = package["klass"]248    args = package["args"]249    kwargs = package["kwargs"]250 251    if strict:252        model = klass(*args, **kwargs)253    else:254        sig = inspect.signature(klass)255        for key in list(kwargs):256            if key not in sig.parameters:257                warnings.warn("Dropping inexistant parameter " + key)258                del kwargs[key]259        model = klass(*args, **kwargs)260 261    state = package["state"]262    training_args = package["training_args"]263    quantizer = get_quantizer(model, training_args)264 265    set_state(model, quantizer, state)266    return model267 268 269def get_state(model, quantizer):270    if quantizer is None:271        state = {k: p.data.to('cpu') for k, p in model.state_dict().items()}272    else:273        state = quantizer.get_quantized_state()274        buf = io.BytesIO()275        th.save(state, buf)276        state = {'compressed': zlib.compress(buf.getvalue())}277    return state278 279 280def set_state(model, quantizer, state):281    if quantizer is None:282        model.load_state_dict(state)283    else:284        buf = io.BytesIO(zlib.decompress(state["compressed"]))285        state = th.load(buf, "cpu")286        quantizer.restore_quantized_state(state)287 288    return state289 290 291def save_state(state, path):292    buf = io.BytesIO()293    th.save(state, buf)294    sig = hashlib.sha256(buf.getvalue()).hexdigest()[:8]295 296    path = path.parent / (path.stem + "-" + sig + path.suffix)297    path.write_bytes(buf.getvalue())298 299 300def save_model(model, quantizer, training_args, path):301    args, kwargs = model._init_args_kwargs302    klass = model.__class__303 304    state = get_state(model, quantizer)305 306    save_to = path307    package = {308        'klass': klass,309        'args': args,310        'kwargs': kwargs,311        'state': state,312        'training_args': training_args,313    }314    th.save(package, save_to)315 316 317def capture_init(init):318    @functools.wraps(init)319    def __init__(self, *args, **kwargs):320        self._init_args_kwargs = (args, kwargs)321        init(self, *args, **kwargs)322 323    return __init__324