mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-05-05 23:22:34 -05:00
6491a7512e
- Removed temporary test files and audit reports - Deleted backup and temp_holding directories - Reorganized module structure (07->09 spatial, 09->07 dataloader) - Added new modules: 11-14 (tokenization, embeddings, attention, transformers) - Updated examples with historical ML milestones - Cleaned up documentation structure
96 lines
2.9 KiB
Python
Generated
96 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, Module # Use the same Module class as layers
|
|
from ..core.spatial import Conv2d
|
|
|
|
# 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 |