Skip to content

Add ring-shared weights and same depth padding #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
YugnatD wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Add ring-shared weights and same depth padding #8

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
202 changes: 173 additions & 29 deletions src/hexagdly.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,6 +340,74 @@ def operation_with_single_hexbase_stride(self, input):
)


# ----------------------------------------------------------------------------
# Ring sharing (share_neighbors): tie weights by hexagonal ring, like TDSCAN.
# The hexagdly offset layout has no clean closed-form hex distance, so the ring
# index of every kernel cell is derived EMPIRICALLY: a single-tap impulse through
# the conv reveals each cell's physical (row, col) offset, and the ring is the
# smallest kernel size whose support contains it. Exact, framework-self-consistent.
# ----------------------------------------------------------------------------

_RING_MAP_CACHE = {}


def _tap_offset(n, i, r, c):
"""Physical (dr, dc) offset of sub-kernel cell (i, r, c) for kernel size n."""
g = 6 * n + 11
cen = g // 2
imp = torch.zeros(1, 1, g, g)
imp[0, 0, cen, cen] = 1.0
sub_kernels = []
for k in range(n + 1):
kh = 2 * n + 1 - k
kw = 1 if k == 0 else 2
a = np.zeros((1, 1, kh, kw), dtype=np.float32)
if k == i:
a[0, 0, r, c] = 1.0
sub_kernels.append(a)
layer = Conv2d_CustomKernel(sub_kernels=sub_kernels, stride=1)
out = layer(imp).detach().numpy()[0, 0]
pos = np.argwhere(np.isclose(out, 1.0))
return int(pos[0][0] - cen), int(pos[0][1] - cen)


def ring_maps_2d(n):
"""Return ``(ring_maps, num_rings)`` for kernel size ``n`` (see hexagdly_tf)."""
if n in _RING_MAP_CACHE:
return _RING_MAP_CACHE[n]
support = {}
for nn in range(1, n + 1):
offs = set()
for i in range(nn + 1):
rows = 2 * nn + 1 - i
cols = 1 if i == 0 else 2
for r in range(rows):
for c in range(cols):
offs.add(_tap_offset(nn, i, r, c))
support[nn] = offs

def ring_of(off):
if off == (0, 0):
return 0
for nn in range(1, n + 1):
if off in support[nn]:
return nn
raise ValueError(f"offset {off} not within kernel size {n}")

ring_maps = []
for i in range(n + 1):
rows = 2 * n + 1 - i
cols = 1 if i == 0 else 2
m = np.zeros((rows, cols), dtype=np.int64)
for r in range(rows):
for c in range(cols):
m[r, c] = ring_of(_tap_offset(n, i, r, c))
ring_maps.append(m)
result = (ring_maps, n + 1)
_RING_MAP_CACHE[n] = result
return result


class Conv2d(HexBase, nn.Module):
r"""Applies a 2D hexagonal convolution`

Expand All @@ -364,7 +432,8 @@ class Conv2d(HexBase, nn.Module):
"""

def __init__(
self, in_channels, out_channels, kernel_size=1, stride=1, bias=True, debug=False
self, in_channels, out_channels, kernel_size=1, stride=1, bias=True,
debug=False, share_neighbors=False
):
super(Conv2d, self).__init__()
self.in_channels = in_channels
Expand All @@ -373,23 +442,32 @@ def __init__(
self.hexbase_stride = stride
self.debug = debug
self.bias = bias
self.share_neighbors = share_neighbors
self.dimensions = 2
self.process = F.conv2d
self.combine = torch.add

for i in range(self.hexbase_size + 1):
setattr(
self,
"kernel" + str(i),
Parameter(
torch.Tensor(
out_channels,
in_channels,
1 + 2 * self.hexbase_size - i,
1 if i == 0 else 2,
)
),
)
if share_neighbors:
# One weight per hex ring; broadcast to every cell at forward time.
self._ring_maps, self.num_rings = ring_maps_2d(self.hexbase_size)
self._ring_idx = [torch.as_tensor(m, dtype=torch.long)
for m in self._ring_maps]
self.ring_weights = Parameter(
torch.Tensor(out_channels, in_channels, self.num_rings))
else:
for i in range(self.hexbase_size + 1):
setattr(
self,
"kernel" + str(i),
Parameter(
torch.Tensor(
out_channels,
in_channels,
1 + 2 * self.hexbase_size - i,
1 if i == 0 else 2,
)
),
)
if self.bias:
self.bias_tensor = Parameter(torch.Tensor(out_channels))
self.kwargs = {"bias": self.bias_tensor}
Expand All @@ -398,6 +476,14 @@ def __init__(
self.init_parameters(self.debug)

def init_parameters(self, debug):
if self.share_neighbors:
if debug:
nn.init.constant_(self.ring_weights, 1)
else:
nn.init.kaiming_normal_(self.ring_weights)
if self.bias:
nn.init.constant_(self.kwargs["bias"], 1.0 if debug else 0.01)
return
if debug:
for i in range(self.hexbase_size + 1):
nn.init.constant_(getattr(self, "kernel" + str(i)), 1)
Expand All @@ -409,7 +495,24 @@ def init_parameters(self, debug):
if self.bias:
nn.init.constant_(getattr(self, "kwargs")["bias"], 0.01)

def _materialize_shared_kernels(self):
"""kernel{i}[:, :, r, c] = ring_weights[:, :, ring_map[i][r, c]].

Gathers along the ring axis -> dense (out, in, rows, cols) kernels, so
the forward pass is unchanged and gradients flow back into ring_weights
(all cells of a ring share one weight), exactly like TDSCAN.
"""
for i in range(self.hexbase_size + 1):
idx = self._ring_idx[i].to(self.ring_weights.device) # (rows, cols)
# ring_weights: (out, in, num_rings) -> index_select on last axis,
# then reshape to (out, in, rows, cols).
flat = torch.index_select(self.ring_weights, 2, idx.reshape(-1))
setattr(self, "kernel" + str(i),
flat.reshape(self.out_channels, self.in_channels, *idx.shape))

def forward(self, input):
if self.share_neighbors:
self._materialize_shared_kernels()
if self.hexbase_stride == 1:
return self.operation_with_single_hexbase_stride(input)
else:
Expand Down Expand Up @@ -596,9 +699,13 @@ class Conv3d(HexBase, nn.Module):
"""

def __init__(
self, in_channels, out_channels, kernel_size=1, stride=1, bias=True, debug=False
self, in_channels, out_channels, kernel_size=1, stride=1, bias=True,
debug=False, share_neighbors=False, depth_padding="valid"
):
super(Conv3d, self).__init__()
if depth_padding not in ("valid", "same"):
raise ValueError("depth_padding must be 'valid' or 'same'.")
self.depth_padding = depth_padding
self.in_channels = in_channels
self.out_channels = out_channels
if isinstance(kernel_size, int):
Expand All @@ -617,24 +724,34 @@ def __init__(
self.depth_stride = stride[0]
self.debug = debug
self.bias = bias
self.share_neighbors = share_neighbors
self.dimensions = 3
self.process = F.conv3d
self.combine = torch.add

for i in range(self.hexbase_size + 1):
setattr(
self,
"kernel" + str(i),
Parameter(
torch.Tensor(
out_channels,
in_channels,
self.depth_size,
1 + 2 * self.hexbase_size - i,
1 if i == 0 else 2,
)
),
)
if share_neighbors:
# Share over the hex axes only; depth (time) stays independent ->
# ring_weights is (out, in, depth, num_rings) (cf. TDSCAN L x rings).
self._ring_maps, self.num_rings = ring_maps_2d(self.hexbase_size)
self._ring_idx = [torch.as_tensor(m, dtype=torch.long)
for m in self._ring_maps]
self.ring_weights = Parameter(
torch.Tensor(out_channels, in_channels, self.depth_size, self.num_rings))
else:
for i in range(self.hexbase_size + 1):
setattr(
self,
"kernel" + str(i),
Parameter(
torch.Tensor(
out_channels,
in_channels,
self.depth_size,
1 + 2 * self.hexbase_size - i,
1 if i == 0 else 2,
)
),
)
if self.bias:
self.bias_tensor = Parameter(torch.Tensor(out_channels))
self.kwargs = {"bias": self.bias_tensor}
Expand All @@ -644,6 +761,14 @@ def __init__(
self.init_parameters(self.debug)

def init_parameters(self, debug):
if self.share_neighbors:
if debug:
nn.init.constant_(self.ring_weights, 1)
else:
nn.init.kaiming_normal_(self.ring_weights)
if self.bias:
nn.init.constant_(self.kwargs["bias"], 1.0 if debug else 0.01)
return
if debug:
for i in range(self.hexbase_size + 1):
nn.init.constant_(getattr(self, "kernel" + str(i)), 1)
Expand All @@ -655,7 +780,26 @@ def init_parameters(self, debug):
if self.bias:
nn.init.constant_(getattr(self, "kwargs")["bias"], 0.01)

def _materialize_shared_kernels(self):
"""kernel{i} = ring_weights gathered along the ring axis into
(out, in, depth, rows, cols); depth left independent."""
for i in range(self.hexbase_size + 1):
idx = self._ring_idx[i].to(self.ring_weights.device) # (rows, cols)
flat = torch.index_select(self.ring_weights, 3, idx.reshape(-1))
setattr(self, "kernel" + str(i),
flat.reshape(self.out_channels, self.in_channels,
self.depth_size, *idx.shape))

def forward(self, input):
if self.share_neighbors:
self._materialize_shared_kernels()
if self.depth_padding == "same":
# Symmetric zero-pad the depth axis (NCDHW -> axis 2) so the temporal
# kernel is centred and output depth == input depth, like TDSCAN.
pad = (self.depth_size - 1) // 2
top = pad
bot = self.depth_size - 1 - pad
input = F.pad(input, [0, 0, 0, 0, top, bot]) # pads last dims; here D
if self.hexbase_stride == 1:
return self.operation_with_single_hexbase_stride(input)
else:
Expand Down