In the document, it mentioned that we can work with PyTorch layer in the loss calculation. I am wondering if we can add the PyTorch layer in the intermediate of the MinkUNet.
For example, I would like to add the self-attention layer to create more robust features. So I first decompose the features in the Sparse tensor, then perform the self-attention on those features. Finally, re-compose features to the sparse tensor and pass it to the decoder of MinkUnet.
However, the performance degrades a lot, I am wondering if this integration is feasible or any other suggestions? Thanks so much!
class MinkUNetBase(ResNetBase):
def __init__(self, in_channels, out_channels, D=3):
ResNetBase.__init__(self, in_channels, out_channels, D)
self.attn_layer = nn.TransformerEncoder()
def encoder(self, x):
attn_feats = self.attn_layer(x.F.unsqueeze(1)) # creat attention cross batch
x = ME.SparseTensor(features=attn_feats.squeeze(1),
coordinate_map_key=x.coordinate_map_key,
coordinate_manager=x.coordinate_manager)
return x
def forward(self, x):
out = self.conv0p1s1(x)
out = self.bn0(out)
out_p1 = self.relu(out)
...
# tensor_stride=16
out = self.conv4p8s2(out_b3p8)
out = self.bn4(out)
out = self.relu(out)
out = self.block4(out)
attn_out = self.encoder(out)
out = self.convtr4p16s2(attn_out)
out = self.bntr4(out)
out = self.relu(out)
...
return self.final(out)
Hi @chrischoy !
In the document, it mentioned that we can work with PyTorch layer in the loss calculation. I am wondering if we can add the PyTorch layer in the intermediate of the MinkUNet.
For example, I would like to add the self-attention layer to create more robust features. So I first decompose the features in the Sparse tensor, then perform the self-attention on those features. Finally, re-compose features to the sparse tensor and pass it to the decoder of MinkUnet.
However, the performance degrades a lot, I am wondering if this integration is feasible or any other suggestions? Thanks so much!