File size: 3,814 Bytes
ccf4794 | 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 | # E-OS v5 — FrameGenerator: Glass-Renderer (~1.6M params).
# Input : hint [B,3,HH,HW] (Wireframe-Layout) + state [B,32]
# Output: Glass-Frame [B,3,2H,2W] sigmoid — lernt Apple-Glas + Wasser-Ripples.
# Voll-konvolutional: Training 128x272 -> App 256x544 (Out 512x1088).
# NPU-freundlich: Conv2d/SiLU/FiLM/Concat/NearestUpsample/PixelShuffle/Sigmoid.
import torch
import torch.nn as nn
import torch.nn.functional as F
class FiLM(nn.Module):
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 FrameGeneratorV5(nn.Module):
def __init__(self, semb=96):
super().__init__()
self.stem = nn.Conv2d(3, 40, 3, 2, 1) # 40 x H/2
self.e1 = nn.Conv2d(40, 64, 3, 2, 1) # 64 x H/4
self.e2 = nn.Conv2d(64, 128, 3, 2, 1) # 128 x H/8
self.b0 = nn.Conv2d(128, 128, 3, 1, 1)
for i in range(7):
setattr(self, "b%d" % (i + 1), nn.Conv2d(128, 128, 3, 1, 1))
self.d3 = nn.Conv2d(128 + 64, 96, 3, 1, 1) # H/4
self.d2 = nn.Conv2d(96 + 40, 56, 3, 1, 1) # H/2
self.d1 = nn.Conv2d(56 + 3, 28, 3, 1, 1) # H
self.o1 = nn.Conv2d(28 + 3, 18, 3, 1, 1)
self.o2 = nn.Conv2d(18, 12, 3, 1, 1)
self.px = nn.PixelShuffle(2)
self.semb = nn.Sequential(nn.Linear(32, 64), nn.SiLU(), nn.Linear(64, semb))
self.film_b = FiLM(128, 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, 8):
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))
return torch.sigmoid(y)
def grad_loss(y, t):
dyx = y[..., :, 1:] - y[..., :, :-1]
dtx = t[..., :, 1:] - t[..., :, :-1]
dyy = y[..., 1:, :] - y[..., :-1, :]
dty = t[..., 1:, :] - t[..., :-1, :]
return (dyx - dtx).abs().mean() + (dyy - dty).abs().mean()
if __name__ == "__main__":
m = FrameGeneratorV5().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, 128, 272), torch.rand(1, 32))
print("train out:", tuple(y.shape))
with torch.no_grad():
y = m(torch.rand(1, 3, 256, 544), torch.rand(1, 32))
print("app out:", tuple(y.shape))
macs = [0]
def hook(mod, inp, out):
macs[0] += mod.weight.numel() * out.shape[-2] * out.shape[-1]
hs = []
for mod in m.modules():
if isinstance(mod, nn.Conv2d):
hs.append(mod.register_forward_hook(hook))
with torch.no_grad():
m(torch.rand(1, 3, 256, 544), torch.rand(1, 32))
for h in hs:
h.remove()
print("app MACs: %.2fG" % (macs[0] / 1e9))
|