team7/talk_with_wind_test_something

1import math2from typing import Optional, Callable3import torch4import torch.nn as nn5from torch import Tensor6 7 8def make_divisible(v: float, divisor: int, min_value: Optional[int] = None) -> int:9    """10    This function is taken from the original tf repo.11    It ensures that all layers have a channel number that is divisible by 812    It can be seen here:13    https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py14    """15    if min_value is None:16        min_value = divisor17    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)18    # Make sure that round down does not go down by more than 10%.19    if new_v < 0.9 * v:20        new_v += divisor21    return new_v22 23 24def cnn_out_size(in_size, padding, dilation, kernel, stride):25    s = in_size + 2 * padding - dilation * (kernel - 1) - 126    return math.floor(s / stride + 1)27 28 29def collapse_dim(x: Tensor, dim: int, mode: str = "pool", pool_fn:  Callable[[Tensor, int], Tensor] = torch.mean,30                 combine_dim: int = None):31    """32    Collapses dimension of multi-dimensional tensor by pooling or combining dimensions33    :param x: input Tensor34    :param dim: dimension to collapse35    :param mode: 'pool' or 'combine'36    :param pool_fn: function to be applied in case of pooling37    :param combine_dim: dimension to join 'dim' to38    :return: collapsed tensor39    """40    if mode == "pool":41        return pool_fn(x, dim)42    elif mode == "combine":43        s = list(x.size())44        s[combine_dim] *= dim45        s[dim] //= dim46        return x.view(s)47 48 49class CollapseDim(nn.Module):50    def __init__(self, dim: int, mode: str = "pool", pool_fn:  Callable[[Tensor, int], Tensor] = torch.mean,51                 combine_dim: int = None):52        super(CollapseDim, self).__init__()53        self.dim = dim54        self.mode = mode55        self.pool_fn = pool_fn56        self.combine_dim = combine_dim57 58    def forward(self, x):59        return collapse_dim(x, dim=self.dim, mode=self.mode, pool_fn=self.pool_fn, combine_dim=self.combine_dim)60