Fixing gradient in sincos positional encoding in monai/networks/blocks/patchembedding.py

monai/networks/blocks/patchembedding.py

@@ -123,6 +123,7 @@ def __init__(
  with torch.no_grad():
  pos_embeddings = build_sincos_position_embedding(grid_size, hidden_size, spatial_dims)
  self.position_embeddings.data.copy_(pos_embeddings.float())
+ self.position_embeddings.requires_grad = False
  else:
  raise ValueError(f"pos_embed_type {self.pos_embed_type} not supported.")
 

monai/networks/blocks/pos_embed_utils.py

@@ -46,7 +46,7 @@ def build_sincos_position_embedding(
  temperature (float): The temperature for the sin-cos position embedding.
 
  Returns:
- pos_embed (nn.Parameter): The sin-cos position embedding as a learnable parameter.
+ pos_embed (nn.Parameter): The sin-cos position embedding as a fixed parameter.
  """
 
  if spatial_dims == 2:

tests/test_patchembedding.py

@@ -93,6 +93,32 @@ def test_shape(self, input_param, input_shape, expected_shape):
  result = net(torch.randn(input_shape))
  self.assertEqual(result.shape, expected_shape)
 
+ def test_sincos_pos_embed(self):
+ net = PatchEmbeddingBlock(
+ in_channels=1,
+ img_size=(32, 32, 32),
+ patch_size=(8, 8, 8),
+ hidden_size=96,
+ num_heads=8,
+ pos_embed_type="sincos",
+ dropout_rate=0.5,
+ )
+
+ self.assertEqual(net.position_embeddings.requires_grad, False)
+
+ def test_learnable_pos_embed(self):
+ net = PatchEmbeddingBlock(
+ in_channels=1,
+ img_size=(32, 32, 32),
+ patch_size=(8, 8, 8),
+ hidden_size=96,
+ num_heads=8,
+ pos_embed_type="learnable",
+ dropout_rate=0.5,
+ )
+
+ self.assertEqual(net.position_embeddings.requires_grad, True)
+
  def test_ill_arg(self):
  with self.assertRaises(ValueError):
  PatchEmbeddingBlock(