diff options
Diffstat (limited to 'doc/CNN_V2.md')
| -rw-r--r-- | doc/CNN_V2.md | 11 |
1 files changed, 7 insertions, 4 deletions
diff --git a/doc/CNN_V2.md b/doc/CNN_V2.md index 4612d7a..6242747 100644 --- a/doc/CNN_V2.md +++ b/doc/CNN_V2.md @@ -214,12 +214,15 @@ def compute_static_features(rgb, depth): ```python class CNNv2(nn.Module): - def __init__(self, kernel_size=3, num_layers=3): + def __init__(self, kernel_sizes, num_layers=3): super().__init__() + if isinstance(kernel_sizes, int): + kernel_sizes = [kernel_sizes] * num_layers + self.kernel_sizes = kernel_sizes self.layers = nn.ModuleList() # All layers: 12D input (4 prev + 8 static) → 4D output - for i in range(num_layers): + for kernel_size in kernel_sizes: self.layers.append( nn.Conv2d(12, 4, kernel_size=kernel_size, padding=kernel_size//2, bias=False) @@ -247,7 +250,7 @@ class CNNv2(nn.Module): ```python # Hyperparameters -kernel_size = 3 # Uniform 3×3 kernels +kernel_sizes = [3, 3, 3] # Per-layer kernel sizes (e.g., [1,3,5]) num_layers = 3 # Number of CNN layers learning_rate = 1e-3 batch_size = 16 @@ -278,7 +281,7 @@ for epoch in range(epochs): torch.save({ 'state_dict': model.state_dict(), # f32 weights 'config': { - 'kernel_size': 3, + 'kernel_sizes': [3, 3, 3], # Per-layer kernel sizes 'num_layers': 3, 'features': ['p0', 'p1', 'p2', 'p3', 'uv.x', 'uv.y', 'sin10_x', 'bias'] }, |