How can I define a cell rather than a layer?
Opened this issue · 2 comments
GELIELEO commented
How can I define a cell rather than a layer?
KhaleelKhan commented
Hello @GELIELEO ,
Defining a Cell instead of a Layer is not possible with the interface we currently provide since this is a optimized implementation. Cell interface is currently under development.
KhaleelKhan commented
update: please check branch feature/egru_cell
. Example usage is given below.
import torch
import evnn_pytorch as evnn
input_size = 128
hidden_size = 16
batch_size = 5
# setting use_custom_cuda=False makes the model use pytorch code instead of EvNN extension,
# has no effect on the step function which is only implemented in pytorch
egru_cell = evnn.EGRU(input_size, hidden_size, batch_first=False, use_custom_cuda=True)
egru_cell.cuda()
timesteps=25
# `x` is a CUDA tensor with shape [N,T,C]
x = torch.rand([timesteps, batch_size, input_size]).cuda()
y = []
y_t = None
h_t = None
for t in range(timesteps):
x_t = x[t]
y_t, h_t, _ = egru_cell.step(x_t, y_t, h_t)
y.append(y_t)
y = torch.stack(y, dim=0)
# Sanity check, some numerical error from float conversion is expected
assert torch.allclose(y, egru_cell(x)[0], atol=5e-4)