File size: 4,095 Bytes
c978bcc | 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 106 | # E-OS v4 — FrameGenerator: FULL-FRAME generative renderer (~0.92M params).
# Input : hint [B,3,192,416] (rough layout, /255) + state [B,32]
# Output: frame [B,3,384,832] in (0,1) via sigmoid (always 2x hint size)
#
# NPU-friendliness (NNAPI): only Conv2d / SiLU / Mul+Add broadcast (FiLM) /
# Concat / ResizeNearest / ResizeBilinear / Sigmoid — no norm layers, no dynamic
# shapes. Fully convolutional -> trains at half res, infers at full res.
import torch
import torch.nn as nn
import torch.nn.functional as F
class FiLM(nn.Module):
"""scale-shift conditioning: y = x*(1+g) + b (g,b from state embedding)"""
def __init__(self, ch, semb):
super().__init__()
self.g = nn.Linear(semb, ch)
self.b = nn.Linear(semb, ch)
def forward(self, x, s):
g = self.g(s).unsqueeze(-1).unsqueeze(-1)
b = self.b(s).unsqueeze(-1).unsqueeze(-1)
return x * (1.0 + g) + b
def up2(x):
return F.interpolate(x, scale_factor=2.0, mode="nearest")
class FrameGenerator(nn.Module):
def __init__(self, semb=96):
super().__init__()
self.stem = nn.Conv2d(3, 32, 3, 2, 1) # 32 x H/2
self.e1 = nn.Conv2d(32, 48, 3, 2, 1) # 48 x H/4
self.e2 = nn.Conv2d(48, 96, 3, 2, 1) # 96 x H/8 (bottleneck res)
self.b0 = nn.Conv2d(96, 96, 3, 1, 1)
for i in range(6):
setattr(self, "b%d" % (i + 1), nn.Conv2d(96, 96, 3, 1, 1))
self.d3 = nn.Conv2d(96 + 48, 96, 3, 1, 1) # +skip e1 -> H/4
self.d2 = nn.Conv2d(96 + 32, 48, 3, 1, 1) # +skip stem -> H/2
self.d1 = nn.Conv2d(48 + 3, 24, 3, 1, 1) # +hint @ H/1
self.o1 = nn.Conv2d(24 + 3, 16, 3, 1, 1) # + bilinear hint @ H/1
self.o2 = nn.Conv2d(16, 12, 3, 1, 1)
self.px = nn.PixelShuffle(2) # DepthToSpace -> 3 x 2H x 2W
self.semb = nn.Sequential(nn.Linear(32, 64), nn.SiLU(), nn.Linear(64, semb))
self.film_b = FiLM(96, semb)
self.film_d = FiLM(96, semb)
for m in self.modules():
if isinstance(m, nn.Conv2d):
nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
if m.bias is not None:
nn.init.zeros_(m.bias)
if isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
nn.init.zeros_(m.bias)
def forward(self, hint, state):
s = self.semb(state)
x0 = F.silu(self.stem(hint))
x1 = F.silu(self.e1(x0))
b = F.silu(self.e2(x1))
b = F.silu(self.b0(b))
b = self.film_b(b, s)
for i in range(1, 7):
b = F.silu(getattr(self, "b%d" % i)(b))
u3 = F.silu(self.d3(torch.cat([up2(b), x1], 1)))
u3 = self.film_d(u3, s)
u2 = F.silu(self.d2(torch.cat([up2(u3), x0], 1)))
u1 = F.silu(self.d1(torch.cat([up2(u2), hint], 1)))
y = F.silu(self.o1(torch.cat([u1, hint], 1)))
y = self.px(self.o2(y))
y = torch.sigmoid(y)
return y
def count_macs(model, hint_hw):
macs = [0]
hooks = []
def hook(mod, inp, out):
macs[0] += mod.weight.numel() * out.shape[-2] * out.shape[-1]
for mod in model.modules():
if isinstance(mod, nn.Conv2d):
hooks.append(mod.register_forward_hook(hook))
h = torch.rand(1, 3, hint_hw[0], hint_hw[1])
s = torch.rand(1, 32)
with torch.no_grad():
model(h, s)
for hk in hooks:
hk.remove()
return macs[0]
if __name__ == "__main__":
m = FrameGenerator()
m.eval()
n = sum(p.numel() for p in m.parameters())
print("params: %.3fM" % (n / 1e6))
with torch.no_grad():
y = m(torch.rand(1, 3, 96, 208), torch.rand(1, 32))
print("train-res out:", tuple(y.shape))
with torch.no_grad():
y2 = m(torch.rand(1, 3, 192, 416), torch.rand(1, 32))
print("app-res out:", tuple(y2.shape))
print("app-res MACs: %.2fG" % (count_macs(m, (192, 416)) / 1e9))
print("train-res MACs: %.2fG" % (count_macs(m, (96, 208)) / 1e9))
|