import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader import torchvision.datasets as datasets import torchvision.transforms as transforms import matplotlib.pyplot as plt from matplotlib.animation import FuncAnimation, PillowWriter from mpl_toolkits.mplot3d import Axes3D import numpy as np from scipy.spatial.transform import Rotation import argparse # Install torchdiffeq if needed try: from torchdiffeq import odeint except ImportError: import subprocess subprocess.run(['pip', 'install', 'torchdiffeq', '-q']) from torchdiffeq import odeint class QuantumVertexGenerator(nn.Module): """ Converts batch of images to 100×12 3D vertex points. Each sample → 12 learnable vertices in 3D space. The vertices represent "quantum states" that encode spatial information. """ def __init__(self, img_size=28, num_vertices=12, batch_size=100): super().__init__() self.img_size = img_size self.num_vertices = num_vertices self.batch_size = batch_size # Learnable base vertices (12 vertices forming initial shape) # Vertices will be rotated/scaled based on image features self.base_vertices = nn.Parameter(torch.randn(num_vertices, 3) * 0.5) # Feature extraction for vertex modulation self.feature_net = nn.Sequential( nn.Conv2d(1, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.MaxPool2d(2), # 14x14 nn.Conv2d(32, 64, 3, padding=1), nn.BatchNorm2d(64), nn.ReLU(), nn.MaxPool2d(2), # 7x7 nn.Conv2d(64, 128, 3, padding=1), nn.BatchNorm2d(128), nn.ReLU(), nn.AdaptiveAvgPool2d((3, 3)), # 3x3 nn.Flatten(), nn.Linear(128 * 9, 256), nn.ReLU(), ) # Rotation and scale predictors for each vertex self.rotation_predictor = nn.Linear(256, num_vertices * 4) # quaternion per vertex self.scale_predictor = nn.Linear(256, num_vertices * 3) # scale per vertex self.translate_predictor = nn.Linear(256, 3) # global translation def forward(self, x): """ Args: x: (batch_size, 1, 28, 28) images Returns: vertices: (batch_size, num_vertices, 3) 3D vertex points """ batch_size = x.shape[0] # Extract features from image features = self.feature_net(x) # (batch, 256) # Predict transformations quaternions = self.rotation_predictor(features) # (batch, num_vertices * 4) scales = torch.sigmoid(self.scale_predictor(features)) # (batch, num_vertices * 3) translation = torch.tanh(self.translate_predictor(features)) # (batch, 3) # Apply transformations to base vertices vertices_list = [] for b in range(batch_size): sample_vertices = [] for v in range(self.num_vertices): # Get base vertex v_base = self.base_vertices[v] # (3,) # Get rotation quaternion for this vertex q = quaternions[b, v*4:(v+1)*4] # (4,) q = q / (q.norm() + 1e-8) # normalize # Apply rotation using Rodrigues formula v_rot = self._rotate_by_quaternion(v_base, q) # Apply scale scale_v = scales[b, v*3:(v+1)*3] # (3,) v_scaled = v_rot * (scale_v + 0.5) sample_vertices.append(v_scaled) sample_vertices = torch.stack(sample_vertices, dim=0) # (num_vertices, 3) sample_vertices = sample_vertices + translation[b] # apply translation vertices_list.append(sample_vertices) vertices = torch.stack(vertices_list, dim=0) # (batch, num_vertices, 3) return vertices def _rotate_by_quaternion(self, v, q): """Rotate vector v by quaternion q.""" # q = [w, x, y, z] w, x, y, z = q[0], q[1], q[2], q[3] # Quaternion rotation formula vx, vy, vz = v[0], v[1], v[2] result_x = 2*( (w*w + x*x - 0.5)*vx + (x*y - w*z)*vy + (x*z + w*y)*vz ) result_y = 2*( (x*y + w*z)*vx + (w*w + y*y - 0.5)*vy + (y*z - w*x)*vz ) result_z = 2*( (x*z - w*y)*vx + (y*z + w*x)*vy + (w*w + z*z - 0.5)*vz ) return torch.stack([result_x, result_y, result_z]) class QuantumVertexODEFunc(nn.Module): """ ODE function that evolves 100×12 vertex points through time. The dynamics model quantum-like interactions between vertices. """ def __init__(self, vertex_dim=3, num_vertices=12, hidden_dim=128): super().__init__() self.vertex_dim = vertex_dim self.num_vertices = num_vertices # Input: all vertex positions + time self.input_dim = num_vertices * vertex_dim + 1 # Shared dynamics network self.dynamics_net = nn.Sequential( nn.Linear(self.input_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.SiLU(), nn.Dropout(0.1), nn.Linear(hidden_dim, hidden_dim), nn.LayerNorm(hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, num_vertices * vertex_dim), ) # Attention mechanism for vertex interactions self.attention = nn.MultiheadAttention(embed_dim=vertex_dim, num_heads=4, batch_first=True) # Learnable edge weights for vertex graph self.edge_mlp = nn.Sequential( nn.Linear(vertex_dim * 2 + 1, 32), nn.ReLU(), nn.Linear(32, vertex_dim), ) def forward(self, t, vertices_flat): """ Args: t: current time (scalar) vertices_flat: (batch, num_vertices * vertex_dim) flattened vertex positions Returns: d_vertices/dt: (batch, num_vertices * vertex_dim) """ batch_size = vertices_flat.shape[0] # Reshape to (batch, num_vertices, vertex_dim) vertices = vertices_flat.view(batch_size, self.num_vertices, self.vertex_dim) # Compute pairwise distances for edge features pairwise_dist = torch.cdist(vertices, vertices) # (batch, num_v, num_v) # Attention-based vertex interactions attn_out, _ = self.attention(vertices, vertices, vertices) # (batch, num_v, dim) # Combine features combined = torch.cat([ vertices.flatten(1), attn_out.flatten(1), t.expand(batch_size, 1) ], dim=-1) # Compute dynamics dynamics = self.dynamics_net(combined) # Add attention contribution dynamics = dynamics + attn_out.flatten(1) * 0.3 return dynamics class Quantum3DClassifier(nn.Module): """ Quantum 3D Vertex Classifier: - 100 sample batch → 100×12 3D vertices - Vertices evolve through ODE - Classification based on evolved states + statistics """ def __init__(self, num_classes=10, num_vertices=12, t_span=(0, 2)): super().__init__() self.num_classes = num_classes self.num_vertices = num_vertices self.t_span = t_span # Vertex generator self.vertex_gen = QuantumVertexGenerator( img_size=28, num_vertices=num_vertices, batch_size=100 ) # ODE function self.ode_func = QuantumVertexODEFunc( vertex_dim=3, num_vertices=num_vertices, hidden_dim=128 ) # Trajectory encoding self.trajectory_encoder = nn.Sequential( nn.Linear(num_vertices * 3 * 3, 256), # initial + final + mean nn.ReLU(), nn.Linear(256, 128), ) # Vertex statistics self.stat_encoder = nn.Sequential( nn.Linear(num_vertices * 3 * 5, 128), # mean, std, max, min, range per vertex nn.ReLU(), nn.Linear(128, 64), ) # Final classifier self.classifier = nn.Sequential( nn.Linear(192 + 50, 256), # trajectory + stats + time embedding nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, num_classes) ) # Time embedding (learned) self.time_embed = nn.Parameter(torch.randn(50)) def compute_trajectory_features(self, trajectory): """ Compute features from ODE trajectory. trajectory: (num_steps, batch, num_vertices * 3) """ # Initial state initial = trajectory[0] # (batch, num_v * 3) # Final state final = trajectory[-1] # (batch, num_v * 3) # Mean over time mean_traj = trajectory.mean(dim=0) # (batch, num_v * 3) # Std over time std_traj = trajectory.std(dim=0) # (batch, num_v * 3) return torch.cat([initial, final, mean_traj], dim=-1) def compute_vertex_statistics(self, vertices): """ Compute statistics from vertex positions. vertices: (batch, num_vertices, 3) """ # Reshape to (batch, num_vertices, 3) flat = vertices.view(vertices.shape[0], -1) # (batch, num_v * 3) # Compute per-vertex statistics mean = vertices.mean(dim=1) # (batch, 3) std = vertices.std(dim=1) # (batch, 3) max_v = vertices.max(dim=1)[0] # (batch, 3) min_v = vertices.min(dim=1)[0] # (batch, 3) # Concatenate all stats = torch.cat([mean, std, max_v, min_v, max_v - min_v], dim=-1) # (batch, 15) # Expand to match vertex dimension stats = stats.unsqueeze(1).expand(-1, self.num_vertices, -1) # (batch, num_v, 15) stats = stats.reshape(stats.shape[0], -1) # (batch, num_v * 15) return stats def forward(self, x): """ Args: x: (batch, 1, 28, 28) MNIST images Returns: logits: (batch, num_classes) """ batch_size = x.shape[0] # Generate 3D vertices for batch vertices = self.vertex_gen(x) # (batch, num_vertices, 3) vertices_flat = vertices.reshape(batch_size, -1) # (batch, num_v * 3) # Time points for ODE t_points = torch.linspace(self.t_span[0], self.t_span[1], 50, device=x.device) # Solve ODE trajectory = odeint( self.ode_func, vertices_flat, t_points, method='rk4' ) # (num_steps, batch, num_v * 3) # Compute trajectory features traj_features = self.compute_trajectory_features(trajectory) # (batch, num_v * 3 * 3) traj_encoded = self.trajectory_encoder(traj_features) # (batch, 128) # Get final vertex positions final_vertices = trajectory[-1].reshape(batch_size, self.num_vertices, 3) # (batch, num_v, 3) # Compute vertex statistics vertex_stats = self.compute_vertex_statistics(final_vertices) # (batch, num_v * 15) stats_encoded = self.stat_encoder(vertex_stats) # (batch, 64) # Combine all features combined = torch.cat([ traj_encoded, stats_encoded, self.time_embed.unsqueeze(0).expand(batch_size, -1) ], dim=-1) # (batch, 192 + 50 = 242) # Classify logits = self.classifier(combined) return logits class Quantum3DClassifierLite(nn.Module): """ Lighter version with more efficient vertex processing. """ def __init__(self, num_classes=10, num_vertices=12, t_span=(0, 1.5), state_dim_per_vertex=18): super().__init__() self.num_classes = num_classes self.num_vertices = num_vertices self.t_span = t_span self.state_dim_per_vertex = state_dim_per_vertex self.state_dim = num_vertices * state_dim_per_vertex self.position_dim = 3 self.position_slice = slice(0, self.position_dim) self.state_groups = ( 'pos', 'vel', 'acc', 'angles', 'angle_vel', 'angle_acc' ) # Efficient encoder self.encoder = nn.Sequential( nn.Conv2d(1, 16, 3, padding=1), nn.BatchNorm2d(16), nn.ReLU(), nn.MaxPool2d(2), nn.Conv2d(16, 32, 3, padding=1), nn.BatchNorm2d(32), nn.ReLU(), nn.AdaptiveAvgPool2d((4, 4)), nn.Flatten(), nn.Linear(32 * 16, 128), nn.ReLU(), ) # Initial latent state per vertex: # 3 pos, 3 vel, 3 acc, 3 angles, 3 angle vel, 3 angle acc. self.vertex_head = nn.Linear(128, self.state_dim) self.vertex_scale = nn.Linear(128, self.state_dim) # Nonlinear ODE in the full latent state space. self.ode_net = nn.Sequential( nn.Linear(self.state_dim + 1, 384), nn.LayerNorm(384), nn.SiLU(), nn.Linear(384, 384), nn.LayerNorm(384), nn.SiLU(), nn.Linear(384, self.state_dim), ) # Classifier self.classifier = nn.Sequential( nn.Linear(self.state_dim * 4, 256), # initial + final + mean + std nn.ReLU(), nn.Dropout(0.3), nn.Linear(256, num_classes) ) def ode_func(self, t, state): return self.ode_net(torch.cat([state, t.expand(state.shape[0], 1)], dim=-1)) def forward(self, x): batch_size = x.shape[0] # Encode image features = self.encoder(x) # (batch, 128) # Generate vertices vertices = self.vertex_head(features) # (batch, num_v * 3) scales = torch.sigmoid(self.vertex_scale(features)) * 2 + 0.5 vertices = vertices * scales # Solve ODE t_points = torch.linspace(self.t_span[0], self.t_span[1], 30, device=x.device) trajectory = odeint(self.ode_func, vertices, t_points, method='rk4') # Compute features from trajectory initial = trajectory[0] final = trajectory[-1] mean_t = trajectory.mean(dim=0) std_t = trajectory.std(dim=0) combined = torch.cat([initial, final, mean_t, std_t], dim=-1) return self.classifier(combined) def get_vertex_trajectory(model, images, num_steps=30): """Return initial vertices and ODE trajectory for either model variant.""" device = images.device t_end = model.t_span[1] if hasattr(model, 't_span') else 1.5 t_points = torch.linspace(0, t_end, num_steps, device=device) if hasattr(model, 'vertex_gen'): vertices = model.vertex_gen(images) vertices_flat = vertices.reshape(images.shape[0], -1) trajectory = odeint(model.ode_func, vertices_flat, t_points, method='rk4') initial_vertices = vertices position_trajectory = trajectory.reshape(num_steps, images.shape[0], model.num_vertices, 3) else: features = model.encoder(images) vertices_flat = model.vertex_head(features) scales = torch.sigmoid(model.vertex_scale(features)) * 2 + 0.5 vertices_flat = vertices_flat * scales trajectory = odeint(model.ode_func, vertices_flat, t_points, method='rk4') latent_vertices = vertices_flat.reshape( images.shape[0], model.num_vertices, model.state_dim_per_vertex ) initial_vertices = latent_vertices[:, :, model.position_slice] position_trajectory = trajectory.reshape( num_steps, images.shape[0], model.num_vertices, model.state_dim_per_vertex )[:, :, :, model.position_slice] return initial_vertices, position_trajectory, trajectory, t_points def visualize_vertices(model, test_loader, device, num_samples=8): """Visualize generated 3D vertices and evolution.""" model.eval() samples = next(iter(test_loader)) images, labels = samples[0][:num_samples].to(device), samples[1][:num_samples] with torch.no_grad(): vertices, position_trajectory, _, t_points = get_vertex_trajectory(model, images, num_steps=30) # Create figure fig = plt.figure(figsize=(16, 12)) # Plot sample images and their vertices for i in range(min(num_samples, 8)): # Image ax1 = fig.add_subplot(4, 6, i*2 + 1) ax1.imshow(images[i].cpu().squeeze(), cmap='gray') ax1.set_title(f'Digit: {labels[i].item()}', fontsize=10) ax1.axis('off') # Initial 3D vertices ax2 = fig.add_subplot(4, 6, i*2 + 2, projection='3d') v_init = vertices[i].cpu().numpy() ax2.scatter(v_init[:, 0], v_init[:, 1], v_init[:, 2], c=range(12), cmap='viridis', s=100) ax2.set_title('Initial Vertices', fontsize=8) ax2.set_xlabel('X') ax2.set_ylabel('Y') ax2.set_zlabel('Z') # Plot evolution of one sample ax3 = fig.add_subplot(4, 6, (13, 14), projection='3d') sample_idx = 0 t_indices = np.linspace(0, position_trajectory.shape[0] - 1, 4, dtype=int) for t_idx in t_indices: v_traj = position_trajectory[t_idx, sample_idx].cpu().numpy() ax3.scatter(v_traj[:, 0], v_traj[:, 1], v_traj[:, 2], label=f't={t_points[t_idx].item():.2f}', s=50, alpha=0.8) ax3.legend() ax3.set_title('Vertex Evolution (Sample 0)') # Trajectory over time ax4 = fig.add_subplot(4, 6, 15) for v in range(0, model.num_vertices, 3): vals = position_trajectory[:, sample_idx, v].cpu() dists = torch.norm(vals, dim=1).numpy() ax4.plot(t_points.cpu().numpy(), dists, label=f'Vertex {v}') ax4.set_xlabel('Time') ax4.set_ylabel('Distance from origin') ax4.set_title('Vertex Distance Over Time') ax4.legend(fontsize=8) ax4.grid(True) # Phase plot (x, y, z for vertex 0) ax5 = fig.add_subplot(4, 6, 16, projection='3d') v0_traj = position_trajectory[:, sample_idx, 0].cpu().numpy() colors = plt.cm.plasma(np.linspace(0, 1, position_trajectory.shape[0])) for i in range(position_trajectory.shape[0] - 1): ax5.plot(v0_traj[i:i+2, 0], v0_traj[i:i+2, 1], v0_traj[i:i+2, 2], color=colors[i], linewidth=2) ax5.scatter(v0_traj[0, 0], v0_traj[0, 1], v0_traj[0, 2], c='green', s=100, label='Start') ax5.scatter(v0_traj[-1, 0], v0_traj[-1, 1], v0_traj[-1, 2], c='red', s=100, label='End') ax5.legend() ax5.set_title('Vertex 0 Trajectory in 3D') plt.tight_layout() plt.savefig('quantum_vertices.png', dpi=150, bbox_inches='tight') plt.show() def animate_vertices(model, test_loader, device, num_samples=1, output_path='quantum_vertices.gif', num_steps=60): """Export an animation of the learned ODE vertex motion.""" model.eval() samples = next(iter(test_loader)) images, labels = samples[0][:num_samples].to(device), samples[1][:num_samples] with torch.no_grad(): logits = model(images) predictions = logits.argmax(dim=1) initial_vertices, position_trajectory, full_trajectory, t_points = get_vertex_trajectory( model, images, num_steps=num_steps ) sample_idx = 0 initial_np = initial_vertices[sample_idx].detach().cpu().numpy() traj_np = position_trajectory[:, sample_idx].detach().cpu().numpy() t_np = t_points.detach().cpu().numpy() if hasattr(model, 'state_dim_per_vertex'): latent_np = full_trajectory.reshape( num_steps, images.shape[0], model.num_vertices, model.state_dim_per_vertex )[:, sample_idx].detach().cpu().numpy() velocity_mag = np.linalg.norm(latent_np[:, :, 3:6], axis=2) else: velocity_mag = np.ones((num_steps, model.num_vertices)) xyz_min = min(initial_np.min(), traj_np.min()) xyz_max = max(initial_np.max(), traj_np.max()) pad = max(0.25, 0.1 * (xyz_max - xyz_min + 1e-6)) fig = plt.figure(figsize=(10, 5)) ax_img = fig.add_subplot(1, 2, 1) ax_3d = fig.add_subplot(1, 2, 2, projection='3d') ax_img.imshow(images[sample_idx].detach().cpu().squeeze(), cmap='gray') ax_img.set_title( f'True: {labels[sample_idx].item()} | Pred: {predictions[sample_idx].item()}', fontsize=11 ) ax_img.axis('off') vertex_colors = plt.cm.viridis(np.linspace(0, 1, model.num_vertices)) trail_lines = [] for color in vertex_colors: line, = ax_3d.plot([], [], [], color=color, linewidth=1.5, alpha=0.6) trail_lines.append(line) scatter = ax_3d.scatter( traj_np[0, :, 0], traj_np[0, :, 1], traj_np[0, :, 2], c=vertex_colors, s=80 ) title = ax_3d.set_title('') ax_3d.set_xlabel('X') ax_3d.set_ylabel('Y') ax_3d.set_zlabel('Z') ax_3d.set_xlim(xyz_min - pad, xyz_max + pad) ax_3d.set_ylim(xyz_min - pad, xyz_max + pad) ax_3d.set_zlim(xyz_min - pad, xyz_max + pad) def update(frame_idx): coords = traj_np[frame_idx] sizes = 40 + 80 * (velocity_mag[frame_idx] / (velocity_mag.max() + 1e-6)) scatter._offsets3d = (coords[:, 0], coords[:, 1], coords[:, 2]) scatter.set_sizes(sizes) for vertex_idx, line in enumerate(trail_lines): history = traj_np[:frame_idx + 1, vertex_idx, :] line.set_data(history[:, 0], history[:, 1]) line.set_3d_properties(history[:, 2]) title.set_text( f'12-Vertex ODE Motion | t={t_np[frame_idx]:.2f} | ' f"state={getattr(model, 'state_dim_per_vertex', 3)}D/vertex" ) return [scatter, title, *trail_lines] anim = FuncAnimation(fig, update, frames=num_steps, interval=90, blit=False) anim.save(output_path, writer=PillowWriter(fps=12)) plt.close(fig) print(f"Animation saved to {output_path}") def visualize_training(train_losses, train_accs, test_accs): """Plot training progress.""" fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) ax1.plot(train_losses, 'b-', linewidth=2) ax1.set_title('Training Loss', fontsize=14) ax1.set_xlabel('Epoch') ax1.set_ylabel('Loss') ax1.grid(True) ax2.plot(train_accs, 'b-', label='Train', linewidth=2) ax2.plot(test_accs, 'r-', label='Test', linewidth=2) ax2.set_title('Accuracy', fontsize=14) ax2.set_xlabel('Epoch') ax2.set_ylabel('Accuracy (%)') ax2.legend() ax2.grid(True) plt.tight_layout() plt.savefig('training_progress.png', dpi=150) plt.show() def train_model(model, train_loader, test_loader, epochs=8, device='cuda'): """Train the Quantum 3D classifier.""" model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs) train_losses = [] train_accs = [] test_accs = [] print(f"Training on {device}") print("=" * 60) for epoch in range(epochs): model.train() running_loss = 0.0 correct = 0 total = 0 for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() running_loss += loss.item() _, predicted = output.max(1) total += target.size(0) correct += predicted.eq(target).sum().item() if batch_idx % 150 == 0: print(f' Epoch {epoch+1} | Batch {batch_idx}/{len(train_loader)} | Loss: {loss.item():.4f}') scheduler.step() # Evaluate train_acc = 100. * correct / total test_acc = evaluate(model, test_loader, device) avg_loss = running_loss / len(train_loader) train_losses.append(avg_loss) train_accs.append(train_acc) test_accs.append(test_acc) print(f'Epoch {epoch+1}/{epochs} | Loss: {avg_loss:.4f} | Train: {train_acc:.2f}% | Test: {test_acc:.2f}%') print('-' * 60) return train_losses, train_accs, test_accs def evaluate(model, test_loader, device='cuda'): """Evaluate model on test set.""" model.eval() correct = 0 total = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) _, predicted = output.max(1) total += target.size(0) correct += predicted.eq(target).sum().item() return 100. * correct / total def save_checkpoint(model, path): """Save the trained model checkpoint.""" checkpoint = { 'model_type': model.__class__.__name__, 'num_classes': model.num_classes, 'num_vertices': model.num_vertices, 't_span': model.t_span, 'state_dim_per_vertex': getattr(model, 'state_dim_per_vertex', 3), 'state_dict': model.state_dict(), } torch.save(checkpoint, path) print(f"Checkpoint saved to {path}") def load_checkpoint(path, device='cpu'): """Load a saved model checkpoint.""" checkpoint = torch.load(path, map_location=device) model_type = checkpoint['model_type'] if model_type == 'Quantum3DClassifierLite': model = Quantum3DClassifierLite( num_classes=checkpoint['num_classes'], num_vertices=checkpoint['num_vertices'], t_span=tuple(checkpoint['t_span']), state_dim_per_vertex=checkpoint.get('state_dim_per_vertex', 3), ) elif model_type == 'Quantum3DClassifier': model = Quantum3DClassifier( num_classes=checkpoint['num_classes'], num_vertices=checkpoint['num_vertices'], t_span=tuple(checkpoint['t_span']), ) else: raise ValueError(f"Unsupported model type in checkpoint: {model_type}") model.load_state_dict(checkpoint['state_dict']) model = model.to(device) model.eval() return model def main(args): # Configuration DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu' BATCH_SIZE = 100 EPOCHS = 1 NUM_VERTICES = 12 STATE_DIM_PER_VERTEX = 18 CHECKPOINT_PATH = 'ex04_model.pt' print("Trainable reference implementation: ex04.py") print(f"Device: {DEVICE}") print(f"Batch Size: {BATCH_SIZE}") print(f"Vertices per sample: {NUM_VERTICES}") print(f"State dim per vertex: {STATE_DIM_PER_VERTEX}") print("=" * 60) # Load MNIST transform = transforms.Compose([ transforms.ToTensor(), ]) train_dataset = datasets.MNIST(root='../data', train=True, download=True, transform=transform) test_dataset = datasets.MNIST(root='../data', train=False, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=2) test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=2) print(f"Train: {len(train_dataset)} | Test: {len(test_dataset)}") print("=" * 60) if args.load_checkpoint: model = load_checkpoint(args.load_checkpoint, DEVICE) print(f"\nLoaded checkpoint: {args.load_checkpoint}") else: # Create model - Use lite version for faster training model = Quantum3DClassifierLite( num_classes=10, num_vertices=NUM_VERTICES, t_span=(0, 1.5), state_dim_per_vertex=STATE_DIM_PER_VERTEX ) print("\nModel Architecture:") print(f"Input: 100 × 1 × 28 × 28 images") print(f"Output: 100 × {NUM_VERTICES} × {model.state_dim_per_vertex} latent state") print("State layout: pos, vel, acc, angles, angle_vel, angle_acc") print(f"ODE t_span: (0, 1.5)") print("=" * 60) if args.load_checkpoint: train_losses, train_accs, test_accs = [], [], [] else: # Train train_losses, train_accs, test_accs = train_model( model, train_loader, test_loader, epochs=EPOCHS, device=DEVICE ) save_checkpoint(model, CHECKPOINT_PATH) # Final results print("\n" + "=" * 60) print("FINAL RESULTS") print("=" * 60) final_test_acc = evaluate(model, test_loader, DEVICE) print(f"Final Test Accuracy: {final_test_acc:.2f}%") print("=" * 60) # Visualize if train_losses: visualize_training(train_losses, train_accs, test_accs) visualize_vertices(model, test_loader, DEVICE) animate_vertices( model, test_loader, DEVICE, num_samples=1, output_path=args.animation_output, num_steps=args.animation_steps ) return model, train_losses, train_accs, test_accs if __name__ == '__main__': parser = argparse.ArgumentParser(description='Train or visualize the ex04 quantum vertex model.') parser.add_argument('--load-checkpoint', type=str, default='', help='Load a saved checkpoint instead of training a new model.') parser.add_argument('--animation-output', type=str, default='quantum_vertices.gif', help='Path for the exported animation.') parser.add_argument('--animation-steps', type=int, default=60, help='Number of ODE frames to render in the animation.') parsed_args = parser.parse_args() model, losses, train_acc, test_acc = main(parsed_args)