Untitled

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import numpy as np
import os

# 1. DEFINE THE MODEL (Matches your Verilog exactly)
# 784 inputs (28x28 pixels), 10 outputs (Digits 0-9)
class SimpleHardwareNet(nn.Module):
    def __init__(self):
        super(SimpleHardwareNet, self).__init__()
        self.fc = nn.Linear(28 * 28, 10) 

    def forward(self, x):
        x = x.view(-1, 28 * 28) # Flatten the 28x28 image
        x = self.fc(x)          # The MAC operation + Bias
        return x

# 2. LOAD DATA AND TRAIN
print("Downloading MNIST and training...")
transform = transforms.Compose([transforms.ToTensor()])
train_data = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_data, batch_size=64, shuffle=True)

model = SimpleHardwareNet()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.005)

# Train for just 2 epochs (Enough for ~90% accuracy on this simple model)
for epoch in range(2):
    for images, labels in train_loader:
        optimizer.zero_grad()
        output = model(images)
        loss = criterion(output, labels)
        loss.backward()
        optimizer.step()
print("Training Complete!")

# 3. EXTRACT AND QUANTIZE WEIGHTS (Float32 -> 8-bit Integer)
print("Extracting and quantizing weights...")
weights = model.fc.weight.data.numpy() # Shape: (10, 784)
biases = model.fc.bias.data.numpy()    # Shape: (10,)

# Find the maximum absolute value to scale everything to fit in 8-bits (-128 to 127)
max_val = max(np.max(np.abs(weights)), np.max(np.abs(biases)))
scale_factor = 127.0 / max_val

# Convert to 8-bit signed integers
quantized_weights = np.round(weights * scale_factor).astype(int)
quantized_biases = np.round(biases * scale_factor).astype(int)

# 4. EXPORT TO VERILOG-READABLE HEX FILES
os.makedirs("verilog_weights", exist_ok=True)

# Helper function to convert integer to 8-bit Two's Complement Hex
def to_hex_8bit(val):
    return f"{(val & 0xFF):02X}"

# Export Weights: One file for each of the 10 neurons
for digit in range(10):
    filename = f"verilog_weights/weights_digit_{digit}.txt"
    with open(filename, "w") as f:
        for w in quantized_weights[digit]:
            f.write(to_hex_8bit(w) + "\n")

# Export Biases: One file containing all 10 biases
with open("verilog_weights/biases.txt", "w") as f:
    for b in quantized_biases:
        f.write(to_hex_8bit(b) + "\n")

print("Export Complete! Check the 'verilog_weights' folder.")