76 lines
2.5 KiB
Python
76 lines
2.5 KiB
Python
"""ScrollCVAE — generates realistic scroll wheel event sequences.
|
||
|
||
Architecture mirrors JointCVAE but smaller (scroll sequences are simpler):
|
||
- Encoder: bidirectional GRU(hidden=64, layers=2)
|
||
- Decoder: unidirectional GRU(hidden=64, layers=2)
|
||
- Input/output: (delta_norm, log_Δt) per time step
|
||
- Condition: [dist_norm, log_dist, direction, viewport_norm, mode_onehot×3] = 7 dims
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
from torch.distributions import Normal
|
||
|
||
|
||
class ScrollCVAE(nn.Module):
|
||
def __init__(
|
||
self,
|
||
seq_len: int = 32,
|
||
latent_dim: int = 16,
|
||
hidden: int = 64,
|
||
cond_dim: int = 7,
|
||
):
|
||
super().__init__()
|
||
self.seq_len = seq_len
|
||
self.latent_dim = latent_dim
|
||
self.hidden = hidden
|
||
self.cond_dim = cond_dim
|
||
self.feat_dim = 2 # (delta_norm, log_Δt)
|
||
|
||
self.enc_gru = nn.GRU(
|
||
input_size=self.feat_dim + cond_dim,
|
||
hidden_size=hidden,
|
||
num_layers=2,
|
||
batch_first=True,
|
||
bidirectional=True,
|
||
)
|
||
self.enc_mu = nn.Linear(hidden * 2, latent_dim)
|
||
self.enc_logvar = nn.Linear(hidden * 2, latent_dim)
|
||
|
||
self.dec_h0 = nn.Linear(latent_dim + cond_dim, hidden * 2)
|
||
self.dec_gru = nn.GRU(
|
||
input_size=latent_dim + cond_dim,
|
||
hidden_size=hidden,
|
||
num_layers=2,
|
||
batch_first=True,
|
||
)
|
||
self.dec_out = nn.Linear(hidden, self.feat_dim)
|
||
|
||
def encode(self, seq: torch.Tensor, cond: torch.Tensor):
|
||
B, T, _ = seq.shape
|
||
c_exp = cond.unsqueeze(1).expand(B, T, self.cond_dim)
|
||
x_in = torch.cat([seq, c_exp], dim=-1)
|
||
_, h_n = self.enc_gru(x_in)
|
||
h_cat = torch.cat([h_n[-2], h_n[-1]], dim=-1)
|
||
return self.enc_mu(h_cat), self.enc_logvar(h_cat)
|
||
|
||
def decode(self, z: torch.Tensor, cond: torch.Tensor):
|
||
B = z.shape[0]
|
||
zc = torch.cat([z, cond], dim=-1)
|
||
h0_flat = self.dec_h0(zc)
|
||
h0 = h0_flat.view(B, 2, self.hidden).permute(1, 0, 2).contiguous()
|
||
inp = zc.unsqueeze(1).expand(B, self.seq_len, -1)
|
||
out, _ = self.dec_gru(inp, h0)
|
||
return self.dec_out(out)
|
||
|
||
def reparameterise(self, mu, logvar):
|
||
std = torch.exp(0.5 * logvar)
|
||
return Normal(mu, std).rsample()
|
||
|
||
def forward(self, seq, cond):
|
||
mu, logvar = self.encode(seq, cond)
|
||
z = self.reparameterise(mu, logvar)
|
||
recon = self.decode(z, cond)
|
||
return recon, mu, logvar
|