Training Code

[Luffy2810]

Hello,

Thank you for sharing this impressive work.

I was wondering if you have any plans to release the training code and model files. If so, could you please share an estimated timeline for when they might be available?

Thank you!

[jdk9405]

Hi, thanks for your interest!
Currently, I have many other commitments, so I’m not able to promise a specific timeline😅
The training code needs to be reimplemented due to some internal issues, but I do plan to release it eventually!

[zhouzq-thu]

Seems this project is from DKM with the following modifications and others:

• pos_conv(in_channels=2) -> ERP_pos_conv(in_channels=3)
• torch.meshgrid: use Pi
• dense_certainty: use 2 channels instead of 1 in terminal_module

Log when export to ONNX

Traceback of TorchScript, serialized code (most recent call last):
  File "code/__torch__/dkm/models/dkm.py", line 94, in match
      dense_certainty1 = torch.sigmoid(dense_certainty0)
      _25 = torch.abs(query_to_support_norm_xy)
      if bool(torch.any(torch.gt(_25, 1))):
         ~~~~ <--- HERE
        _26 = torch.abs(query_to_support_norm_xy)
        _27 = torch.sum(torch.gt(_26, 1), [-1])

Traceback of TorchScript, original code (most recent call last):
  File "/home/dongki/workspace/temp/EDM/dkm/models/dkm.py", line 1093, in match
    
            dense_certainty = dense_certainty.sigmoid()  # logits -> probs
            if (query_to_support_norm_xy.abs() > 1).any() and True:
                ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
                wrong = (query_to_support_norm_xy.abs() > 1).sum(dim=-1) > 0
                dense_certainty[wrong[:,None]] = 0

Difference in state dict:

        size mismatch for decoder.embedding_decoder.terminal_module.32.weight: copying a param with shape torch.Size([4, 384, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 384, 1, 1]).
        size mismatch for decoder.embedding_decoder.terminal_module.32.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([3]).
        size mismatch for decoder.embedding_decoder.terminal_module.16.weight: copying a param with shape torch.Size([4, 384, 1, 1]) from checkpoint, the shape in current model is torch.Size([3, 384, 1, 1]).
        size mismatch for decoder.embedding_decoder.terminal_module.16.bias: copying a param with shape torch.Size([4]) from checkpoint, the shape in current model is torch.Size([3]).
        size mismatch for decoder.gps.32.pos_conv.weight: copying a param with shape torch.Size([256, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([256, 2, 1, 1]).
        size mismatch for decoder.gps.16.pos_conv.weight: copying a param with shape torch.Size([256, 3, 1, 1]) from checkpoint, the shape in current model is torch.Size([256, 2, 1, 1]).

[abrahamezzeddine]

Hello @jdk9405

Any update on when you plan to release the training code? Thank you.