plot.py

import matplotlib.pyplot as plt
train_losses = [
    12.224, 9.330, 8.299, 7.563, 7.139,
     6.636, 6.272, 5.938, 5.601, 5.326,
     5.198, 4.902, 4.711, 4.734, 4.452,
     4.349, 4.124, 3.984, 4.003, 3.795
]

test_losses = [
    9.732, 8.237, 7.064, 5.996, 6.110,
    5.372, 5.255, 5.060, 4.413, 4.096,
    3.975, 3.800, 4.064, 3.818, 3.499,
    3.766, 3.119, 2.889, 3.185, 3.591
]

epochs = range(1, 21)

# Plot 
plt.figure(figsize=(9, 5))

plt.plot(epochs, train_losses, label='Train Loss',
         color='blue', marker='o', markersize=4, linewidth=2)

plt.plot(epochs, test_losses, label='Test Loss',
         color='orange', marker='s', markersize=4, linewidth=2)

# Mark best test result
best_test = min(test_losses)
best_epoch = test_losses.index(best_test) + 1
plt.scatter([best_epoch], [best_test], color='red', zorder=5, s=80)
plt.xlabel('Epoch', fontsize=12)
plt.ylabel('Bits per Dimension', fontsize=12)
plt.title('PixelCNN++ Training Loss\n'
          '(CIFAR-10: 10,000 train / 2,000 test, 64 filters, 20 epochs)',
          fontsize=12)
plt.legend(fontsize=10)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.savefig('loss_curve.png', dpi=150)
plt.show()
print(f"Saved loss_curve.png")
print(f"Best Test Loss: {best_test:.3f} bits/dim at epoch {best_epoch}")