36 lines
1.1 KiB
Python
36 lines
1.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
from . import VarLenTensor
|
|
|
|
__all__ = [
|
|
'SparseLinear'
|
|
]
|
|
|
|
|
|
def chunked_apply(module, x: torch.Tensor, chunk_size: int) -> torch.Tensor:
|
|
if chunk_size <= 0 or x.shape[0] <= chunk_size:
|
|
return module(x)
|
|
|
|
# Process first chunk to determine output shape and dtype
|
|
out_0 = module(x[0:chunk_size])
|
|
out_shape = (x.shape[0],) + out_0.shape[1:]
|
|
out = torch.empty(out_shape, device=x.device, dtype=out_0.dtype)
|
|
out[0:chunk_size] = out_0
|
|
|
|
# Process remaining chunks
|
|
for i in range(chunk_size, x.shape[0], chunk_size):
|
|
out[i:i+chunk_size] = module(x[i:i+chunk_size])
|
|
return out
|
|
|
|
|
|
class SparseLinear(nn.Linear):
|
|
def __init__(self, in_features, out_features, bias=True):
|
|
super(SparseLinear, self).__init__(in_features, out_features, bias)
|
|
self.low_vram = False
|
|
self.chunk_size = 65536
|
|
|
|
def forward(self, input: VarLenTensor) -> VarLenTensor:
|
|
if self.low_vram:
|
|
return input.replace(chunked_apply(super().forward, input.feats, self.chunk_size))
|
|
return input.replace(super().forward(input.feats))
|