Spaces:
Running
Running
File size: 2,951 Bytes
c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 358ab8f c6cae91 |
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 94 95 96 97 98 99 100 101 102 103 104 105 |
import torch
import torch.nn as nn
import torchvision.models as tvm
class Decoder(nn.Module):
def __init__(
self, layers, *args, super_resolution=False, num_prototypes=1, **kwargs
) -> None:
super().__init__(*args, **kwargs)
self.layers = layers
self.scales = self.layers.keys()
self.super_resolution = super_resolution
self.num_prototypes = num_prototypes
def forward(self, features, context=None, scale=None):
if context is not None:
features = torch.cat((features, context), dim=1)
stuff = self.layers[scale](features)
logits, context = (
stuff[:, : self.num_prototypes],
stuff[:, self.num_prototypes :],
)
return logits, context
class ConvRefiner(nn.Module):
def __init__(
self,
in_dim=6,
hidden_dim=16,
out_dim=2,
dw=True,
kernel_size=5,
hidden_blocks=5,
amp=True,
residual=False,
amp_dtype=torch.float16,
):
super().__init__()
self.block1 = self.create_block(
in_dim,
hidden_dim,
dw=False,
kernel_size=1,
)
self.hidden_blocks = nn.Sequential(
*[
self.create_block(
hidden_dim,
hidden_dim,
dw=dw,
kernel_size=kernel_size,
)
for hb in range(hidden_blocks)
]
)
self.hidden_blocks = self.hidden_blocks
self.out_conv = nn.Conv2d(hidden_dim, out_dim, 1, 1, 0)
self.amp = amp
self.amp_dtype = amp_dtype
self.residual = residual
def create_block(
self,
in_dim,
out_dim,
dw=True,
kernel_size=5,
bias=True,
norm_type=nn.BatchNorm2d,
):
num_groups = 1 if not dw else in_dim
if dw:
assert (
out_dim % in_dim == 0
), "outdim must be divisible by indim for depthwise"
conv1 = nn.Conv2d(
in_dim,
out_dim,
kernel_size=kernel_size,
stride=1,
padding=kernel_size // 2,
groups=num_groups,
bias=bias,
)
norm = (
norm_type(out_dim)
if norm_type is nn.BatchNorm2d
else norm_type(num_channels=out_dim)
)
relu = nn.ReLU(inplace=True)
conv2 = nn.Conv2d(out_dim, out_dim, 1, 1, 0)
return nn.Sequential(conv1, norm, relu, conv2)
def forward(self, feats):
b, c, hs, ws = feats.shape
with torch.autocast("cuda", enabled=self.amp, dtype=self.amp_dtype):
x0 = self.block1(feats)
x = self.hidden_blocks(x0)
if self.residual:
x = (x + x0) / 1.4
x = self.out_conv(x)
return x
|