- 移除 JEPA/lejepa-identifiability 子模块 gitlink - 移除 research/multiply/MultiPLY 子模块 gitlink - 删除 .gitmodules（不再有外部 URL 依赖） - 两个目录内容作为普通文件纳入主仓库追踪 - 删除各自内部 .git 目录，消除嵌套 git 仓库
This commit is contained in:
gaojie
@@ -0,0 +1,36 @@
|
||||
"""Data generation: latent sources and OU augmentation."""
|
||||
|
||||
import math
|
||||
import torch
|
||||
|
||||
|
||||
def _gennorm_unit_var_scale(alpha):
|
||||
"""Scale β so gennorm(α, β) has unit variance: β = sqrt(Γ(1/α) / Γ(3/α))."""
|
||||
return math.exp(0.5 * (math.lgamma(1.0 / alpha) - math.lgamma(3.0 / alpha)))
|
||||
|
||||
|
||||
def sample_latents(D, N, dist="gaussian", device="cuda", alpha=None):
|
||||
"""Sample D points in R^N (unit variance)."""
|
||||
if dist == "gaussian":
|
||||
return torch.randn(D, N, device=device)
|
||||
elif dist == "laplace":
|
||||
return torch.distributions.Laplace(0, 1 / (2 ** 0.5)).sample((D, N)).to(device)
|
||||
elif dist == "gennorm":
|
||||
if alpha is None:
|
||||
raise ValueError("gennorm requires alpha")
|
||||
scale = _gennorm_unit_var_scale(alpha)
|
||||
u = torch.distributions.Gamma(1.0 / alpha, 1.0).sample((D, N)).to(device)
|
||||
sign = torch.randint(0, 2, (D, N), device=device).float() * 2 - 1
|
||||
return scale * sign * u.pow(1.0 / alpha)
|
||||
else:
|
||||
raise ValueError(f"Unknown distribution: {dist}")
|
||||
|
||||
|
||||
def ou_augment(z, rho, n_views=2, dist="gaussian", alpha=None):
|
||||
"""OU channel: z' = ρz + √(1-ρ²)η, η drawn from same dist as source.
|
||||
Returns (V, B, N)."""
|
||||
fac = (1 - rho ** 2) ** 0.5
|
||||
D, N = z.shape
|
||||
eta = sample_latents(n_views * D, N, dist=dist, device=z.device, alpha=alpha)
|
||||
eta = eta.reshape(n_views, D, N)
|
||||
return rho * z.unsqueeze(0) + fac * eta
|
||||
@@ -0,0 +1,143 @@
|
||||
"""
|
||||
Core training engine — single function used by all experiments.
|
||||
|
||||
Handles: LR schedule (warmup + cosine), online data generation,
|
||||
loss computation (lejepa or whiten), periodic evaluation of ALL metrics
|
||||
on a fixed eval set, standardized output schema.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from .losses import SIGReg, whitening_loss, alignment_loss, infonce_loss
|
||||
from .data import sample_latents, ou_augment
|
||||
from .metrics import compute_all_metrics
|
||||
|
||||
|
||||
def warmup_cosine_lr(step, total_steps, base_lr):
|
||||
"""Constant for first half, cosine decay for second half."""
|
||||
warmup = total_steps // 2
|
||||
if step < warmup:
|
||||
return base_lr
|
||||
t = (step - warmup) / (total_steps - warmup)
|
||||
return base_lr * 0.5 * (1 + np.cos(np.pi * t))
|
||||
|
||||
|
||||
def train_and_evaluate(
|
||||
encoder,
|
||||
mix_fn,
|
||||
*,
|
||||
N,
|
||||
rho,
|
||||
lamb,
|
||||
sigma=1.0,
|
||||
mode="lejepa",
|
||||
source_dist="gaussian",
|
||||
source_alpha=None,
|
||||
steps=20000,
|
||||
batch_size=256,
|
||||
lr=3e-3,
|
||||
z_eval,
|
||||
log_every=100,
|
||||
device="cuda",
|
||||
):
|
||||
"""Train encoder and evaluate periodically.
|
||||
|
||||
Args:
|
||||
encoder: nn.Module, x -> h
|
||||
mix_fn: callable, z -> x
|
||||
N: latent dimension
|
||||
rho: OU correlation
|
||||
lamb: regularization weight
|
||||
mode: "lejepa" or "whiten"
|
||||
source_dist: "gaussian", "laplace", or "gennorm"
|
||||
steps: total training steps
|
||||
batch_size: batch size (online data)
|
||||
lr: peak learning rate
|
||||
z_eval: (num_eval, N) fixed eval tensor
|
||||
log_every: eval frequency
|
||||
device: torch device string
|
||||
|
||||
Returns:
|
||||
encoder: trained encoder
|
||||
log: dict of lists — training curves and periodic eval metrics
|
||||
"""
|
||||
sigreg = SIGReg().to(device)
|
||||
opt = torch.optim.AdamW(encoder.parameters(), lr=lr)
|
||||
|
||||
# Precompute eval mixing (constant across training)
|
||||
x_eval = mix_fn(z_eval)
|
||||
|
||||
log_keys = [
|
||||
"step", "lr",
|
||||
# training losses
|
||||
"align", "sigreg", "whiten", "total",
|
||||
# eval metrics
|
||||
"r2_zx", "r2_xz", "r2_zh", "r2_hz",
|
||||
"orth_err", "orth_err_normalized",
|
||||
"epsilon", "delta", "D_bound", "approx_bound",
|
||||
"procrustes_mse", "L_h", "trace_cov",
|
||||
]
|
||||
log = {k: [] for k in log_keys}
|
||||
|
||||
for step in range(steps + 1):
|
||||
# LR schedule
|
||||
current_lr = warmup_cosine_lr(step, steps, lr)
|
||||
for pg in opt.param_groups:
|
||||
pg["lr"] = current_lr
|
||||
|
||||
# Online data
|
||||
z_batch = sample_latents(batch_size, N, dist=source_dist,
|
||||
device=device, alpha=source_alpha)
|
||||
z_aug = ou_augment(z_batch, rho, dist=source_dist, alpha=source_alpha) # (2, B, N)
|
||||
h = encoder(mix_fn(z_aug).flatten(0, 1)).reshape(2, batch_size, N)
|
||||
|
||||
align = alignment_loss(h)
|
||||
sig = sigreg(h)
|
||||
wht = whitening_loss(h)
|
||||
|
||||
if mode == "lejepa":
|
||||
loss = lamb * sig + (1 - lamb) * align
|
||||
elif mode == "whiten":
|
||||
loss = lamb * wht + (1 - lamb) * align
|
||||
elif mode == "infonce":
|
||||
loss = infonce_loss(h, sigma)
|
||||
else:
|
||||
raise ValueError(f"Unknown mode: {mode}")
|
||||
|
||||
opt.zero_grad()
|
||||
loss.backward()
|
||||
opt.step()
|
||||
|
||||
if step % log_every == 0 or (step < 1000 and step % 100 == 0):
|
||||
log["step"].append(step)
|
||||
log["lr"].append(current_lr)
|
||||
log["align"].append(align.item())
|
||||
log["sigreg"].append(sig.item())
|
||||
log["whiten"].append(wht.item())
|
||||
log["total"].append(loss.item())
|
||||
|
||||
# Full eval on fixed set
|
||||
encoder.eval()
|
||||
with torch.no_grad():
|
||||
h_eval = encoder(x_eval)
|
||||
z_prime = ou_augment(
|
||||
z_eval, rho, n_views=1,
|
||||
dist=source_dist, alpha=source_alpha
|
||||
).squeeze(0)
|
||||
h_prime = encoder(mix_fn(z_prime))
|
||||
|
||||
metrics = compute_all_metrics(z_eval, x_eval, h_eval, h_prime, rho, N)
|
||||
|
||||
for k, v in metrics.items():
|
||||
log[k].append(v)
|
||||
|
||||
encoder.train()
|
||||
|
||||
if step % (log_every * 10) == 0:
|
||||
print(f" step {step:5d} | lr={current_lr:.1e} "
|
||||
f"align={align.item():.2e} sig={sig.item():.1f} "
|
||||
f"R²(h->z)={metrics['r2_hz']:.4f} "
|
||||
f"orth={metrics['orth_err']:.4f}")
|
||||
|
||||
return encoder, log
|
||||
@@ -0,0 +1,53 @@
|
||||
"""Loss functions."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class SIGReg(nn.Module):
|
||||
"""Sliced characteristic function regularizer (Balestriero & LeCun 2025)."""
|
||||
|
||||
def __init__(self, knots=17, n_slices=256, t_max=3.0):
|
||||
super().__init__()
|
||||
self.n_slices = n_slices
|
||||
t = torch.linspace(0, t_max, knots)
|
||||
dt = t_max / (knots - 1)
|
||||
w = torch.full((knots,), 2 * dt)
|
||||
w[[0, -1]] = dt
|
||||
self.register_buffer("t", t)
|
||||
self.register_buffer("phi", torch.exp(-t**2 / 2))
|
||||
self.register_buffer("weights", w * torch.exp(-t**2 / 2))
|
||||
|
||||
def forward(self, h):
|
||||
"""h: (V, B, N) -> scalar."""
|
||||
flat = h.flatten(0, 1)
|
||||
A = F.normalize(torch.randn(flat.size(-1), self.n_slices, device=flat.device), dim=0)
|
||||
xt = (flat @ A).unsqueeze(-1) * self.t
|
||||
err = (xt.cos().mean(0) - self.phi) ** 2 + xt.sin().mean(0) ** 2
|
||||
return (err @ self.weights).mean() * flat.size(0)
|
||||
|
||||
|
||||
def whitening_loss(h):
|
||||
"""||Cov(h) - I||²_F. h: (V, B, N) -> scalar."""
|
||||
flat = h.flatten(0, 1)
|
||||
flat = flat - flat.mean(dim=0)
|
||||
cov = (flat.T @ flat) / (flat.shape[0] - 1)
|
||||
return (cov - torch.eye(flat.shape[1], device=h.device)).square().mean()
|
||||
|
||||
|
||||
def alignment_loss(h):
|
||||
"""Pull positive-pair views together. h: (V, B, N) -> scalar."""
|
||||
return (h.mean(0) - h).square().mean()
|
||||
|
||||
|
||||
def infonce_loss(h, sigma):
|
||||
"""Symmetric Gaussian-kernel InfoNCE: sim(u, v) = -||u - v||² / (2σ²).
|
||||
h: (V, B, N) with V=2 views. Negatives are other batch elements.
|
||||
"""
|
||||
h1, h2 = h[0], h[1] # (B, N) each
|
||||
d12 = ((h1.unsqueeze(1) - h2.unsqueeze(0)) ** 2).sum(-1) # (B, B)
|
||||
sim = -d12 / (2 * sigma ** 2)
|
||||
loss_a = -(sim.diag() - torch.logsumexp(sim, dim=1)).mean()
|
||||
loss_b = -(sim.diag() - torch.logsumexp(sim, dim=0)).mean()
|
||||
return 0.5 * (loss_a + loss_b)
|
||||
@@ -0,0 +1,65 @@
|
||||
"""Evaluation metrics — standardized across all experiments."""
|
||||
|
||||
import torch
|
||||
|
||||
def bidirectional_r2(a, b):
|
||||
"""R²(a->b) and R²(b->a) via torch lstsq on GPU. a, b are tensors."""
|
||||
def _r2(x, y):
|
||||
x1 = torch.cat([x, torch.ones(len(x), 1, device=x.device)], dim=1)
|
||||
W = torch.linalg.lstsq(x1, y).solution
|
||||
ss_res = ((y - x1 @ W) ** 2).sum()
|
||||
ss_tot = ((y - y.mean(0)) ** 2).sum()
|
||||
return (1 - ss_res / ss_tot).item()
|
||||
return _r2(a, b), _r2(b, a)
|
||||
|
||||
|
||||
def compute_all_metrics(z, x, h, h_prime, rho, N):
|
||||
"""All metrics on GPU. z, x, h, h_prime are torch tensors."""
|
||||
r2_zx, r2_xz = bidirectional_r2(z, x)
|
||||
r2_zh, r2_hz = bidirectional_r2(z, h)
|
||||
|
||||
# Orthogonality
|
||||
z1 = torch.cat([z, torch.ones(len(z), 1, device=z.device)], dim=1)
|
||||
W = torch.linalg.lstsq(z1, h).solution
|
||||
A = W[:N].T
|
||||
orth_err = torch.linalg.norm(A.T @ A - torch.eye(N, device=A.device), 'fro').item()
|
||||
orth_err_normalized = orth_err / (N ** 0.5)
|
||||
|
||||
# Bound quantities
|
||||
cov_h = torch.cov(h.T)
|
||||
epsilon = torch.linalg.norm(cov_h - torch.eye(N, device=h.device), 'fro').item()
|
||||
trace_cov = torch.trace(cov_h).item()
|
||||
L_h = ((h_prime - h) ** 2).sum(dim=1).mean().item()
|
||||
delta = max(L_h - 2 * (1 - rho) * trace_cov, 0.0)
|
||||
spectral_gap = 2 * rho * (1 - rho)
|
||||
D_bound = delta / spectral_gap if spectral_gap > 0 else float("inf")
|
||||
approx_bound = D_bound + (epsilon + D_bound) ** 2
|
||||
|
||||
# Procrustes
|
||||
M = (h.T @ z) / len(z)
|
||||
U, S, Vt = torch.linalg.svd(M)
|
||||
Q = U @ Vt
|
||||
procrustes_mse = ((h - z @ Q.T) ** 2).sum(dim=1).mean().item()
|
||||
|
||||
return {
|
||||
"r2_zx": r2_zx, "r2_xz": r2_xz,
|
||||
"r2_zh": r2_zh, "r2_hz": r2_hz,
|
||||
"orth_err": orth_err, "orth_err_normalized": orth_err_normalized,
|
||||
"epsilon": epsilon, "delta": delta, "D_bound": D_bound,
|
||||
"approx_bound": approx_bound, "procrustes_mse": procrustes_mse,
|
||||
"L_h": L_h, "trace_cov": trace_cov,
|
||||
}
|
||||
|
||||
|
||||
def compute_recovery_metrics(z, h, N, suffix=""):
|
||||
"""R² in both directions + orthogonality. Suffix appended to keys."""
|
||||
r2_zh, r2_hz = bidirectional_r2(z, h)
|
||||
z1 = torch.cat([z, torch.ones(len(z), 1, device=z.device)], dim=1)
|
||||
W = torch.linalg.lstsq(z1, h).solution
|
||||
A = W[:N].T
|
||||
orth_err = torch.linalg.norm(A.T @ A - torch.eye(N, device=A.device), 'fro').item()
|
||||
return {
|
||||
f"r2_zh{suffix}": r2_zh, f"r2_hz{suffix}": r2_hz,
|
||||
f"orth_err{suffix}": orth_err,
|
||||
f"orth_err_normalized{suffix}": orth_err / (N ** 0.5),
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Nonlinear mixing functions."""
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
# ── 2D mixing functions ──────────────────────────────────────────────────────
|
||||
|
||||
def mix_spiral(z):
|
||||
"""g(z) = R(π‖z‖) z — measure-preserving spiral diffeomorphism."""
|
||||
norms = z.norm(dim=-1) * torch.pi
|
||||
c, s = norms.cos(), norms.sin()
|
||||
R = torch.stack([torch.stack([c, -s], dim=-1),
|
||||
torch.stack([s, c], dim=-1)], dim=-2)
|
||||
return (R @ z.unsqueeze(-1)).squeeze(-1)
|
||||
|
||||
|
||||
def mix_banana(z):
|
||||
"""Banana: x0 = z0, x1 = z1 + z0²."""
|
||||
return torch.stack([z[..., 0], z[..., 1] + z[..., 0] ** 2], dim=-1)
|
||||
|
||||
|
||||
def mix_sinusoid(z):
|
||||
"""Sinusoidal shear: x0 = z0 + sin(1.5 z1), x1 = z1."""
|
||||
return torch.stack([z[..., 0] + torch.sin(1.5 * z[..., 1]), z[..., 1]], dim=-1)
|
||||
|
||||
|
||||
MIXINGS_2D = {
|
||||
"spiral": mix_spiral,
|
||||
"banana": mix_banana,
|
||||
"sinusoid": mix_sinusoid,
|
||||
# "nvp" handled via make_coupling_mixing(N=2, n_layers=...)
|
||||
}
|
||||
|
||||
|
||||
# ── Coupling-layer mixing (any dimension) ────────────────────────────────────
|
||||
|
||||
def make_coupling_mixing(N, n_layers=4, seed=1337, device="cuda"):
|
||||
"""RealNVP-style coupling layers. Works for any even N (including N=2)."""
|
||||
half = N // 2
|
||||
torch.manual_seed(seed)
|
||||
Ws = []
|
||||
for _ in range(n_layers):
|
||||
W, _ = torch.linalg.qr(torch.randn(half, half, device=device))
|
||||
Ws.append(W * 2.0)
|
||||
|
||||
def mix(z):
|
||||
for i, W in enumerate(Ws):
|
||||
z1, z2 = z[..., :half], z[..., half:]
|
||||
if i % 2 == 0:
|
||||
z2 = z2 + torch.tanh(z1 @ W)
|
||||
else:
|
||||
z1 = z1 + torch.tanh(z2 @ W)
|
||||
z = torch.cat([z1, z2], dim=-1)
|
||||
return z
|
||||
|
||||
return mix
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Encoder architectures."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import numpy as np
|
||||
|
||||
|
||||
def make_mlp_encoder(N, hidden=256, n_layers=4, device="cuda"):
|
||||
"""MLP encoder."""
|
||||
layers = [nn.Linear(N, hidden), nn.GELU()]
|
||||
for _ in range(n_layers - 1):
|
||||
layers += [nn.Linear(hidden, hidden), nn.GELU()]
|
||||
layers.append(nn.Linear(hidden, N))
|
||||
return nn.Sequential(*layers).to(device)
|
||||
|
||||
|
||||
class MatchedEncoder(nn.Module):
|
||||
"""Inverse coupling-layer encoder matched to NVP mixing architecture."""
|
||||
|
||||
def __init__(self, N, n_layers=4, device="cuda"):
|
||||
super().__init__()
|
||||
half = N // 2
|
||||
self.half = half
|
||||
self.n_layers = n_layers
|
||||
self.Ws = nn.ParameterList([
|
||||
nn.Parameter(torch.randn(half, half, device=device) / np.sqrt(half))
|
||||
for _ in range(n_layers)
|
||||
])
|
||||
|
||||
def forward(self, x):
|
||||
for i, W in reversed(list(enumerate(self.Ws))):
|
||||
z1, z2 = x[..., :self.half], x[..., self.half:]
|
||||
if i % 2 == 0:
|
||||
z2 = z2 - torch.tanh(z1 @ W)
|
||||
else:
|
||||
z1 = z1 - torch.tanh(z2 @ W)
|
||||
x = torch.cat([z1, z2], dim=-1)
|
||||
return x
|
||||
|
||||
|
||||
def make_matched_encoder(N, n_layers=4, seed=42, device="cuda"):
|
||||
torch.manual_seed(seed)
|
||||
return MatchedEncoder(N, n_layers=n_layers, device=device).to(device)
|
||||
|
||||
|
||||
def make_cnn_encoder(d_latent=2, device="cuda"):
|
||||
return nn.Sequential(
|
||||
nn.Conv2d(3, 32, 4, 2, 1), nn.BatchNorm2d(32), nn.GELU(),
|
||||
nn.Conv2d(32, 64, 4, 2, 1), nn.BatchNorm2d(64), nn.GELU(),
|
||||
nn.Conv2d(64, 128, 4, 2, 1), nn.BatchNorm2d(128), nn.GELU(),
|
||||
nn.Conv2d(128, 256, 4, 2, 1), nn.BatchNorm2d(256), nn.GELU(),
|
||||
torch.nn.AvgPool2d(4), nn.Flatten(),
|
||||
nn.Linear(256, 256), nn.BatchNorm1d(256), nn.GELU(),
|
||||
nn.Linear(256, d_latent),
|
||||
).to(device)
|
||||
@@ -0,0 +1,95 @@
|
||||
"""DMC Reacher rendering and dataset utilities."""
|
||||
|
||||
import os
|
||||
os.environ.setdefault("MUJOCO_GL", "egl")
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from dm_control import suite
|
||||
from tqdm import tqdm
|
||||
|
||||
from .data import ou_augment
|
||||
|
||||
|
||||
def make_env():
|
||||
return suite.load(domain_name="reacher", task_name="hard")
|
||||
|
||||
|
||||
def render_at(env, qpos, target, height=64, width=64):
|
||||
"""Set joint angles and render → (3, H, W) float32 in [0, 1]."""
|
||||
env.physics.data.qpos[:2] = qpos
|
||||
env.physics.data.qvel[:] = 0
|
||||
env.physics.named.model.geom_pos['target', :2] = target
|
||||
env.physics.forward()
|
||||
rgb = env.physics.render(height=height, width=width, camera_id=0)
|
||||
return rgb.transpose(2, 0, 1).astype(np.float32) / 255.0
|
||||
|
||||
|
||||
def render_batch(env, qpos_batch, target, height=64, width=64):
|
||||
"""Render batch → (N, 3, H, W)."""
|
||||
N = len(qpos_batch)
|
||||
imgs = np.empty((N, 3, height, width), dtype=np.float32)
|
||||
for i in tqdm(range(N), desc="Rendering"):
|
||||
imgs[i] = render_at(env, qpos_batch[i], target, height, width)
|
||||
return imgs
|
||||
|
||||
|
||||
def generate_ou_image_pairs(env, N, rho, target, seed=9999):
|
||||
"""
|
||||
Sample OU latent pairs, render both → (img_t, img_tp1, z_t, z_tp1).
|
||||
|
||||
Uses the same OU process as the rest of the repo but renders through MuJoCo.
|
||||
"""
|
||||
rng = np.random.default_rng(seed)
|
||||
z_t = rng.standard_normal((N, 2)).astype(np.float32)
|
||||
eps = rng.standard_normal((N, 2)).astype(np.float32)
|
||||
z_tp1 = rho * z_t + np.sqrt(1 - rho**2) * eps
|
||||
|
||||
print(f"Rendering {2 * N} images (rho={rho})...")
|
||||
img_t = render_batch(env, z_t, target)
|
||||
img_tp1 = render_batch(env, z_tp1, target)
|
||||
return img_t, img_tp1, z_t, z_tp1
|
||||
|
||||
|
||||
def normalize_images(img_t, img_tp1, img_eval=None):
|
||||
"""Per-channel mean/std normalization. Returns normalized arrays + stats."""
|
||||
mean = img_t.mean(axis=(0, 2, 3), keepdims=True)
|
||||
std = img_t.std(axis=(0, 2, 3), keepdims=True) + 1e-6
|
||||
img_t = (img_t - mean) / std
|
||||
img_tp1 = (img_tp1 - mean) / std
|
||||
if img_eval is not None:
|
||||
img_eval = (img_eval - mean) / std
|
||||
return img_t, img_tp1, img_eval, mean, std
|
||||
return img_t, img_tp1, mean, std
|
||||
|
||||
|
||||
def solve_ik_grid(env, target, n_grid=200):
|
||||
"""Find joint angles that place fingertip at target via grid search."""
|
||||
best_dist, best_qpos = np.inf, None
|
||||
for q0 in np.linspace(-np.pi, np.pi, n_grid):
|
||||
for q1 in np.linspace(-np.pi, np.pi, n_grid):
|
||||
env.physics.data.qpos[:2] = [q0, q1]
|
||||
env.physics.named.model.geom_pos['target', :2] = target
|
||||
env.physics.forward()
|
||||
tip = env.physics.named.data.geom_xpos['finger'][:2]
|
||||
d = np.linalg.norm(tip - target)
|
||||
if d < best_dist:
|
||||
best_dist = d
|
||||
best_qpos = np.array([q0, q1])
|
||||
return best_qpos, best_dist
|
||||
|
||||
|
||||
class ReacherOUDataset(torch.utils.data.Dataset):
|
||||
"""Prerendered OU image pairs with ground-truth latents."""
|
||||
|
||||
def __init__(self, img_t, img_tp1, z_t, z_tp1):
|
||||
self.img_t = torch.from_numpy(img_t)
|
||||
self.img_tp1 = torch.from_numpy(img_tp1)
|
||||
self.z_t = torch.from_numpy(z_t)
|
||||
self.z_tp1 = torch.from_numpy(z_tp1)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.img_t)
|
||||
|
||||
def __getitem__(self, i):
|
||||
return self.img_t[i], self.img_tp1[i], self.z_t[i], self.z_tp1[i]
|
||||
Reference in New Issue
Block a user