mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-06-04 04:27:22 -05:00
e6d0757bbd
Import Strategy: - Keep explicit 'from tinytorch.core.spatial import Conv2d' - Maps directly to module structure (Module 09 → core.spatial) - Better for education: students see exactly where each concept lives - Removed redundant tinytorch/nn.py (nn/ directory already exists) Milestone 04 Optimizations: - Reduced epochs: 50 → 20 (explicit loops are slow!) - Print progress every 5 epochs (instead of 10) - Load from local npz file (no sklearn dependency) - Still achieves ~80%+ accuracy Educational Rationale: TinyTorch uses explicit imports to show module structure: tinytorch.core.tensor # Module 01 tinytorch.core.layers # Module 03 tinytorch.core.spatial # Module 09 tinytorch.core.losses # Module 04 PyTorch's torch.nn is convenient but pedagogically unclear. Our approach: clarity over convenience!
97 lines
2.9 KiB
Python
Generated
97 lines
2.9 KiB
Python
Generated
"""
|
|
TinyTorch Neural Network Module (nn)
|
|
|
|
This package provides PyTorch-compatible neural network building blocks:
|
|
|
|
Core Components:
|
|
- Module: Base class for all layers (automatic parameter registration)
|
|
- Linear: Fully connected layer (renamed from Dense)
|
|
- Conv2d: 2D convolutional layer (renamed from MultiChannelConv2D)
|
|
|
|
Functional Interface:
|
|
- functional (F): Stateless operations like relu, flatten, max_pool2d
|
|
|
|
Example Usage:
|
|
import tinytorch.nn as nn
|
|
import tinytorch.nn.functional as F
|
|
|
|
class CNN(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.conv1 = nn.Conv2d(3, 32, (3, 3)) # RGB → 32 features
|
|
self.fc1 = nn.Linear(800, 10) # Classifier
|
|
|
|
def forward(self, x):
|
|
x = F.relu(self.conv1(x)) # Convolution + activation
|
|
x = F.flatten(x) # Flatten for dense layer
|
|
return self.fc1(x) # Classification
|
|
|
|
model = CNN()
|
|
params = list(model.parameters()) # Auto-collected parameters!
|
|
|
|
The key insight: Students implement the core algorithms (conv, linear transforms)
|
|
while this infrastructure provides the clean API they expect from PyTorch.
|
|
"""
|
|
|
|
# Import layers from core (these contain the student implementations)
|
|
from ..core.layers import Linear, ReLU, Dropout
|
|
from ..core.activations import Sigmoid
|
|
from ..core.spatial import Conv2d, MaxPool2d, AvgPool2d
|
|
|
|
# Import transformer components
|
|
from ..core.embeddings import Embedding, PositionalEncoding
|
|
from ..core.attention import SelfAttention, scaled_dot_product_attention
|
|
from ..core.transformers import LayerNorm, TransformerBlock
|
|
|
|
# Import functional interface
|
|
from . import functional
|
|
|
|
# Make functional available as F (PyTorch convention)
|
|
import tinytorch.nn.functional as F
|
|
|
|
# Utility functions
|
|
def Parameter(data, requires_grad=True):
|
|
"""Create a parameter tensor (learnable weight)."""
|
|
from ..core.tensor import Tensor
|
|
import numpy as np
|
|
if not isinstance(data, Tensor):
|
|
data = Tensor(np.array(data))
|
|
data.requires_grad = requires_grad
|
|
return data
|
|
|
|
class Sequential(Module):
|
|
"""Sequential container for stacking layers."""
|
|
def __init__(self, *layers):
|
|
super().__init__()
|
|
self.layers = layers
|
|
|
|
def forward(self, x):
|
|
for layer in self.layers:
|
|
x = layer(x) if hasattr(layer, '__call__') else layer.forward(x)
|
|
return x
|
|
|
|
def parameters(self):
|
|
params = []
|
|
for layer in self.layers:
|
|
if hasattr(layer, 'parameters'):
|
|
params.extend(layer.parameters())
|
|
return params
|
|
|
|
# Export the main public API
|
|
__all__ = [
|
|
'Module',
|
|
'Linear',
|
|
'Conv2d',
|
|
'Embedding',
|
|
'PositionalEncoding',
|
|
'SelfAttention',
|
|
'LayerNorm',
|
|
'TransformerBlock',
|
|
'Sequential',
|
|
'Parameter',
|
|
'scaled_dot_product_attention',
|
|
'functional',
|
|
'F'
|
|
]
|
|
|
|
# Note: Parameter function will be available after tensor module export |