scratch.py

import os
import torch
import torch.nn as nn
import matplotlib.pyplot as plt

# 1. Setup Layer and Input
conv = nn.Conv2d(in_channels=1, out_channels=16, kernel_size=3, stride=2, padding=1)
fake_image = torch.randn(1, 1, 28, 28)

# 2. Forward Pass
output = conv(fake_image)

# 3. Visualization
fig, axes = plt.subplots(1, 2, figsize=(10, 5))

# Plot Original "Image" (28x28)
# We use .squeeze() to remove dimensions of size 1 so matplotlib can read it
axes[0].imshow(fake_image.squeeze().detach().numpy(), cmap='gray')
axes[0].set_title(f"Input Shape: {fake_image.shape}\n(28x28)")

# Plot one of the 16 Output Feature Maps (14x14)
# We pick index [0, 0, :, :] (First batch, first filter)
axes[1].imshow(output[0, 0].detach().numpy(), cmap='viridis')
axes[1].set_title(f"Output Shape: {output.shape}\n(14x14 - Filter #1)")

os.makedirs("outputs", exist_ok=True)
plt.tight_layout()
plt.savefig("outputs/shape_vis.png")
print("Visualization saved as 'outputs/shape_vis.png'. Open it in Cursor to view!")

# Access the actual weights of the first filter
weights = conv.weight[0, 0] # Get the 3x3 grid for the 1st filter
bias = conv.bias[0]

print("The 3x3 Kernel Weights for Filter #1:")
print(weights)
print(f"The Bias for Filter #1: {bias.item()}")