|
import torch |
|
|
|
|
|
class InPlaceSetSlice(torch.autograd.Function): |
|
|
|
@staticmethod |
|
def forward(ctx, full_tensor, last_slice, x_idx, x_val): |
|
full_tensor[x_idx] = x_val |
|
ctx.x_idx = x_idx |
|
ret = torch.Tensor().to(device=full_tensor.device, dtype=torch.bfloat16) |
|
ret.set_(full_tensor[:x_idx + 1]) |
|
return ret |
|
|
|
@staticmethod |
|
def backward(ctx, grad_out): |
|
if ctx.x_idx == 0: |
|
return None, None, None, grad_out[ctx.x_idx] |
|
else: |
|
return None, grad_out[:ctx.x_idx], None, grad_out[ctx.x_idx] |
|
|
|
|
|
def apply_inplace_set(x_acc, x_idx, x_val): |
|
full_tensor, last_slice = x_acc |
|
new_slice = InPlaceSetSlice.apply(full_tensor, last_slice, x_idx, x_val) |
|
return full_tensor, new_slice |
|
|
|
|
|
class DWAModules(torch.nn.Module): |
|
|
|
def __init__(self, n_blocks, dilation=1, period=1): |
|
super().__init__() |
|
self.n_blocks = n_blocks |
|
self.dilation = dilation |
|
self.period = period |
|
self.alphas = torch.nn.ModuleList([torch.nn.Linear((i+1+dilation)//dilation, 1, bias=False) if (i+1)%period == 0 else None for i in range(n_blocks)]) |
|
self.accumulators = None |
|
self._init_weights() |
|
|
|
def _init_weights(self): |
|
for module in self.alphas: |
|
if module is not None: |
|
module.weight.data.zero_() |
|
module.weight.data[0, -1] = 1. |
|
|
|
def init_accumulators(self, x): |
|
x_accs = [] |
|
for i in range(self.dilation): |
|
current_group_size = (self.n_blocks + 1) // self.dilation |
|
if i < (self.n_blocks + 1) % self.dilation: |
|
current_group_size += 1 |
|
x_accs.append((torch.zeros((current_group_size, *x.shape), device=x.device, dtype=x.dtype), None)) |
|
x_accs[0] = apply_inplace_set(x_accs[0], 0, x) |
|
self.accumulators = x_accs |
|
|
|
def forward(self, x, block_idx): |
|
assert self.accumulators is not None, "`init_accumulators(x)` needs to be called first" |
|
self.accumulators[(block_idx+1) % self.dilation] = apply_inplace_set( |
|
self.accumulators[(block_idx+1) % self.dilation], |
|
(block_idx+1)//self.dilation, |
|
x |
|
) |
|
if (block_idx+1) % self.period == 0: |
|
x = torch.tensordot(self.alphas[block_idx].weight.view(-1), self.accumulators[(block_idx+1)%self.dilation][1], dims=1) |
|
return x |
|
|
