Spaces:
Sleeping
Sleeping
File size: 2,804 Bytes
9b0d6c2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 |
import math
from typing import Optional, Callable
import torch
import torch.nn as nn
from torch import Tensor
def NAME_TO_WIDTH(name):
frame_mn_map = {
'frame_mn01': 0.1,
'frame_mn02': 0.2,
'frame_mn04': 0.4,
'frame_mn05': 0.5,
'frame_mn06': 0.6,
'frame_mn08': 0.8,
'frame_mn10': 1.0,
'frame_mn12': 1.2,
'frame_mn14': 1.4,
'frame_mn16': 1.6,
'frame_mn20': 2.0,
'frame_mn30': 3.0,
'frame_mn40': 4.0,
}
frame_dymn_map = {
'frame_dymn04': 0.4,
'frame_dymn10': 1.0,
'frame_dymn20': 2.0,
}
try:
if name.startswith('frame_dymn'):
w = frame_dymn_map[name[:len('frame_dymnxx')]]
else:
w = frame_mn_map[name[:len('frame_mnxx')]]
except:
w = 1.0
return w
def make_divisible(v: float, divisor: int, min_value: Optional[int] = None) -> int:
"""
This function is taken from the original tf repo.
It ensures that all layers have a channel number that is divisible by 8
It can be seen here:
https://github.com/tensorflow/models/blob/master/research/slim/nets/mobilenet/mobilenet.py
"""
if min_value is None:
min_value = divisor
new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
# Make sure that round down does not go down by more than 10%.
if new_v < 0.9 * v:
new_v += divisor
return new_v
def cnn_out_size(in_size, padding, dilation, kernel, stride):
s = in_size + 2 * padding - dilation * (kernel - 1) - 1
return math.floor(s / stride + 1)
def collapse_dim(x: Tensor, dim: int, mode: str = "pool", pool_fn: Callable[[Tensor, int], Tensor] = torch.mean,
combine_dim: int = None):
"""
Collapses dimension of multi-dimensional tensor by pooling or combining dimensions
:param x: input Tensor
:param dim: dimension to collapse
:param mode: 'pool' or 'combine'
:param pool_fn: function to be applied in case of pooling
:param combine_dim: dimension to join 'dim' to
:return: collapsed tensor
"""
if mode == "pool":
return pool_fn(x, dim)
elif mode == "combine":
s = list(x.size())
s[combine_dim] *= dim
s[dim] //= dim
return x.view(s)
class CollapseDim(nn.Module):
def __init__(self, dim: int, mode: str = "pool", pool_fn: Callable[[Tensor, int], Tensor] = torch.mean,
combine_dim: int = None):
super(CollapseDim, self).__init__()
self.dim = dim
self.mode = mode
self.pool_fn = pool_fn
self.combine_dim = combine_dim
def forward(self, x):
return collapse_dim(x, dim=self.dim, mode=self.mode, pool_fn=self.pool_fn, combine_dim=self.combine_dim) |