scripts/train_toy_stg.py
与本次数学页面一起发布的源码快照。SHA-256: e838a776c07d595b58a9860601af4711464b425e83cb3037d33f1051af2b02e2
1"""Small, synthetic inverse-rendering lesson. NOT the official STG trainer.
2
3Fit mean, linear motion, opacity and RGB to synthetic multiview images using
4autograd, export the exact STG-Lite record, then render it with native_smoke.
5No SfM, densification, learned rotation, feature decoder or real-world claims.
6Requires torch and numpy; uses CUDA when present. No downloaded training data.
7"""
8import argparse
9import json
10import math
11from pathlib import Path
12import time
13
14import numpy as np
15import torch
16
17PROPERTIES = ("x y z trbf_center trbf_scale nx ny nz "
18 "motion_0 motion_1 motion_2 motion_3 motion_4 motion_5 motion_6 motion_7 motion_8 "
19 "f_dc_0 f_dc_1 f_dc_2 opacity scale_0 scale_1 scale_2 "
20 "rot_0 rot_1 rot_2 rot_3 omega_0 omega_1 omega_2 omega_3").split()
21
22
23def render(mean, motion, rgb_logits, opacity_logits, times, cameras, width=64, height=48):
24 # OpenCV camera axes; all cameras point along +Z. Focal length scales with width.
25 device = mean.device
26 focal = width * 60 / 64
27 dt = times[:, None, None] - 0.5
28 xyz = mean[None] + motion[None] * dt - cameras[:, None, :]
29 x, y, z = xyz.unbind(-1)
30 uv = torch.stack((focal*x/z + width/2, focal*y/z + height/2), -1)
31 zero = torch.zeros_like(z)
32 jacobian = torch.stack((focal/z, zero, -focal*x/z.square(),
33 zero, focal/z, -focal*y/z.square()), -1).reshape(*z.shape, 2, 3)
34 covariance = 0.12**2 * (jacobian @ jacobian.transpose(-1, -2))
35 covariance = covariance + torch.eye(2, device=device) * 0.3
36 ys, xs = torch.meshgrid(torch.arange(height, device=device)+0.5,
37 torch.arange(width, device=device)+0.5, indexing="ij")
38 delta = torch.stack((xs, ys), -1)[None, None] - uv[:, :, None, None]
39 inverse = torch.linalg.inv(covariance)
40 power = -0.5 * torch.einsum("bnhwi,bnij,bnhwj->bnhw", delta, inverse, delta)
41 temporal = torch.exp(-((times[:, None]-0.5)/0.55).square())
42 alpha = (torch.sigmoid(opacity_logits)[None, :, None, None]
43 * temporal[:, :, None, None] * torch.exp(power)).clamp(max=0.99)
44 alpha = torch.where(alpha >= 1/255, alpha, 0)
45 order = z.argsort(dim=1) # Discrete ordering; no gradient through visibility changes.
46 alpha = alpha.gather(1, order[:, :, None, None].expand_as(alpha))
47 color = torch.sigmoid(rgb_logits)[order]
48 transmittance = torch.cumprod(torch.cat((torch.ones_like(alpha[:, :1]), 1-alpha), 1), 1)[:, :-1]
49 return (transmittance[:, :, :, :, None] * alpha[:, :, :, :, None]
50 * color[:, :, None, None, :]).sum(1)
51
52
53def main():
54 parser = argparse.ArgumentParser(description=__doc__)
55 parser.add_argument("--output", type=Path, default=Path("artifacts/toy-training"))
56 parser.add_argument("--steps", type=int, default=400)
57 parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda" if torch.cuda.is_available() else "cpu")
58 args = parser.parse_args()
59 if args.steps < 1:
60 parser.error("--steps must be positive")
61 torch.manual_seed(7)
62 args.output.mkdir(parents=True, exist_ok=True)
63 device = torch.device(args.device)
64 teacher = [torch.tensor(value, device=device, dtype=torch.float32) for value in [
65 [[-.55,-.25,3.0],[.0,-.25,3.2],[.55,-.25,3.4],[-.3,.3,3.1],[.3,.3,3.3]],
66 [[.4,.1,0],[-.2,.3,.1],[.1,-.4,0],[.3,-.1,-.1],[-.4,-.1,.1]],
67 [[2,-1,-1],[-1,2,-1],[-1,-1,2],[2,1,-1],[-1,1,2]],
68 [2,2,2,2,2],
69 ]]
70 times = torch.linspace(0.05, 0.95, 9, device=device).repeat(2)
71 cameras = torch.zeros((18, 3), device=device)
72 cameras[:9, 0] = -0.3
73 cameras[9:, 0] = 0.3
74 with torch.no_grad():
75 truth = render(*teacher, times, cameras)
76 fitted = [torch.nn.Parameter(teacher[0]+torch.randn_like(teacher[0])*0.08),
77 torch.nn.Parameter(torch.zeros_like(teacher[1])),
78 torch.nn.Parameter(teacher[2]+torch.randn_like(teacher[2])*0.5),
79 torch.nn.Parameter(torch.ones_like(teacher[3]))]
80 optimizer = torch.optim.Adam(fitted, lr=0.025)
81 history = []
82 started = time.perf_counter()
83 for step in range(args.steps):
84 optimizer.zero_grad(set_to_none=True)
85 loss = (render(*fitted, times, cameras)-truth).square().mean()
86 if not torch.isfinite(loss):
87 raise RuntimeError("Non-finite training loss")
88 loss.backward()
89 optimizer.step()
90 history.append(float(loss.detach().cpu()))
91 if step % 100 == 0:
92 print(f"step={step} mse={history[-1]:.8f}", flush=True)
93 with torch.no_grad():
94 final_mse = float((render(*fitted, times, cameras)-truth).square().mean().cpu())
95 held_times = torch.tensor([0.35,0.65], device=device)
96 held_cameras = torch.zeros((2,3), device=device)
97 held_mse = float((render(*fitted, held_times, held_cameras)-render(*teacher, held_times, held_cameras)).square().mean().cpu())
98 rows = np.zeros((len(teacher[0]),32), dtype="<f4")
99 rows[:, :3] = fitted[0].detach().cpu().numpy()
100 rows[:, 3] = 0.5
101 rows[:, 4] = math.log(0.55)
102 rows[:, 8:11] = fitted[1].detach().cpu().numpy()
103 rows[:, 17:20] = fitted[2].sigmoid().detach().cpu().numpy()
104 rows[:, 20] = fitted[3].detach().cpu().numpy()
105 rows[:, 21:24] = math.log(0.12)
106 rows[:, 24] = 1 # quaternion wxyz
107 header = (f"ply\nformat binary_little_endian 1.0\nelement vertex {len(rows)}\n"
108 + "".join(f"property float {p}\n" for p in PROPERTIES) + "end_header\n")
109 (args.output/"model.ply").write_bytes(header.encode("ascii")+rows.tobytes())
110 (args.output/"cameras.json").write_text(json.dumps({
111 "schema":"pajama.stg.reference-cameras.v1", "cameras":[{
112 "name":"synthetic-heldout", "width":64, "height":48, "fx":60, "fy":60,
113 "position":[0,0,0], "rotation":[[1,0,0],[0,1,0],[0,0,1]]}]}, indent=2), encoding="utf-8")
114 report = {
115 "scope":"synthetic teaching experiment; not official STG training or real-scene reconstruction",
116 "torch":torch.__version__, "device":str(device), "seed":7, "steps":args.steps,
117 "splats":len(rows), "trainImages":18, "imageSize":[64,48],
118 "learned":["mean","linear motion","RGB","opacity"],
119 "fixed":["scale","rotation","temporal center","temporal scale","splat count"],
120 "initialMse":history[0], "finalMse":final_mse,
121 "heldoutSyntheticMse":held_mse, "heldoutSyntheticPsnr":-10*math.log10(max(held_mse,1e-15)),
122 "elapsedSeconds":time.perf_counter()-started,
123 "timeContract":{"normalizedStart":0,"normalizedEnd":1,"captureSeconds":None},
124 "lossEvery100Steps":history[::100],
125 }
126 (args.output/"training-report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
127 print(json.dumps(report, indent=2))
128 if final_mse >= history[0] * 0.1:
129 raise RuntimeError("Training failed to reduce loss by at least 10x")
130
131
132if __name__ == "__main__":
133 main()
134