81 lines
2.5 KiB
Python
81 lines
2.5 KiB
Python
import torch
|
|
import torch.optim as optim
|
|
import matplotlib.pyplot as plt
|
|
from hsmm_model import GaussianHSMM
|
|
from hsmm_inference import viterbi_decode
|
|
from toy_data import generate_toy_data
|
|
|
|
# --- Settings ---
|
|
N_STATES = 10
|
|
INPUT_DIM = 5 # Matches the 'dim' in generate_toy_data
|
|
MAX_DUR = 50
|
|
LR = 0.05
|
|
EPOCHS = 20
|
|
|
|
# In train():
|
|
def train():
|
|
print("1. Generating Data...")
|
|
train_data = generate_toy_data(n_samples=30, seq_len=300, n_clusters=N_STATES, dim=INPUT_DIM)
|
|
|
|
print("2. Initializing Model...")
|
|
model = GaussianHSMM(N_STATES, INPUT_DIM, MAX_DUR)
|
|
optimizer = optim.Adam(model.parameters(), lr=LR)
|
|
|
|
print("3. Training Loop...")
|
|
loss_history = []
|
|
|
|
for epoch in range(EPOCHS):
|
|
epoch_loss = 0
|
|
optimizer.zero_grad()
|
|
|
|
# Batching: Gradient Accumulation
|
|
for seq in train_data:
|
|
loss = model(seq) # Forward pass
|
|
loss.backward() # Backward pass
|
|
epoch_loss += loss.item()
|
|
|
|
# Normalize gradients
|
|
for p in model.parameters():
|
|
if p.grad is not None:
|
|
p.grad /= len(train_data)
|
|
|
|
optimizer.step()
|
|
loss_history.append(epoch_loss)
|
|
|
|
if epoch % 5 == 0:
|
|
print(f"Epoch {epoch:02d} | NLL Loss: {epoch_loss:.2f}")
|
|
|
|
# --- Verification ---
|
|
print("\n4. Results:")
|
|
learned_means = model.means.detach().view(-1).numpy()
|
|
learned_means.sort()
|
|
print(f"True Means: [-5.0, 0.0, 5.0]")
|
|
print(f"Learned Means: {learned_means}")
|
|
|
|
# --- Visualization Block in main.py ---
|
|
print("5. Visualizing Inference...")
|
|
test_seq = train_data[0]
|
|
predicted_path = viterbi_decode(model, test_seq)
|
|
|
|
fig, ax = plt.subplots(2, 1, figsize=(12, 6), sharex=True)
|
|
|
|
# Plot 1: The Multi-Dimensional Data (Transposed so Time is X-axis)
|
|
# This shows the "features" changing color as the state changes
|
|
ax[0].imshow(test_seq.numpy().T, aspect='auto', cmap='viridis', interpolation='nearest')
|
|
ax[0].set_title(f"Raw Data ({INPUT_DIM} Dimensions)")
|
|
ax[0].set_ylabel("Feature Dim")
|
|
|
|
# Plot 2: The Inferred States
|
|
# Reshape path to (1, T) for imshow
|
|
path_img = np.array(predicted_path)[np.newaxis, :]
|
|
ax[1].imshow(path_img, aspect='auto', cmap='tab10', interpolation='nearest')
|
|
ax[1].set_title("Inferred HSMM States")
|
|
ax[1].set_ylabel("State ID")
|
|
ax[1].set_xlabel("Time (Frames)")
|
|
|
|
plt.tight_layout()
|
|
plt.show()
|
|
|
|
if __name__ == "__main__":
|
|
train()
|