def subtract_layer(x, device):
y = torch.zeros(x.shape, x.shape)
y = y.to(device)
for i in range(x.shape):
if i % 2 == 0:
y[i, :] = x[i, :] - x[i + 1, :]
y[i + 1, :] = x[i + 1, :] - x[i, :]
In this layer, x’s shape is (batchsize, channels), i need the feature compare result of two neighbor samples.
Is the code correct? Or should I rewrite this layer with forward and backward codes?
As a function, you don’t need to care about the BP.
Additionally, it will be more formal to write a subclass of nn.Module if there are some member variables or more complex logic in it.
I rewrite as follow, are they the right code format?
def init(self, device):
self.device = device
def forward(self, x): output = torch.FloatTensor(x.shape, x.shape).to(self.device) for i in range(x.shape): if i % 2 == 0: output[i, :] = x[i, :] - x[i + 1, :] output[i + 1, :] = x[i + 1, :] - x[i, :] return output
I think it can achieve your intention.