dmlc / dmlc/dgl

dgl.nn.NodeEmbedding.all_set_embedding not work in cpu device

Open
#7,897 1 comment 0 reactions 0 assignees View on GitHub
stale-issue
Dominant language
Python
Stars
14.3k
Forks
3.1k
PR merge metrics
No merged PRs in 30d

Description

## 🐛 Bug

dgl.nn.NodeEmbedding.all_set_embedding not work in cpu device.

## To Reproduce

```python
import torch
from dgl.nn import NodeEmbedding

emb_num = 3
emb_size = 10
ckpt_path = "model.pth"

def save(device):
emb_layer = NodeEmbedding(emb_num, emb_size, "emb", device=device,
init_func=torch.nn.init.xavier_uniform_)
init_emb = emb_layer.all_get_embedding()
print("init:\n", init_emb)
torch.save({"emb": init_emb}, ckpt_path)

def load(device):
emb_layer = NodeEmbedding(emb_num, emb_size, "emb", device=device)
print("rank:", emb_layer._rank)

checkpoint = torch.load(ckpt_path, torch.device(device))
print("ckpt:\n", checkpoint["emb"])

emb_layer.all_set_embedding(checkpoint["emb"])
print("run all_set_embedding ...")

aget_emb = emb_layer.all_get_embedding()
print("aget:\n", aget_emb)

if __name__ == "__main__":
save("cpu")
load("cpu")
```

output:

```bash
init:
tensor([[-0.0215, -0.2327, -0.2921, 0.3921, -0.0462, 0.2429, -0.3395, -0.3044,
-0.6371, 0.5789],
[-0.6599, 0.1928, 0.2490, -0.2308, -0.6601, 0.5161, -0.5522, 0.1596,
0.5544, 0.0445],
[ 0.4633, 0.2364, 0.2244, -0.1086, -0.6691, 0.2808, -0.0796, 0.0587,
0.4433, -0.3488]])
rank: -1
ckpt:
tensor([[-0.0215, -0.2327, -0.2921, 0.3921, -0.0462, 0.2429, -0.3395, -0.3044,
-0.6371, 0.5789],
[-0.6599, 0.1928, 0.2490, -0.2308, -0.6601, 0.5161, -0.5522, 0.1596,
0.5544, 0.0445],
[ 0.4633, 0.2364, 0.2244, -0.1086, -0.6691, 0.2808, -0.0796, 0.0587,
0.4433, -0.3488]])
run all_set_embedding ...
aget:
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]])
```

## Environment

- DGL Version (e.g., 1.0): 2.4.0
- Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): 2.3.1
- OS (e.g., Linux): Linux
- How you installed DGL (`conda`, `pip`, source): pip
- Build command you used (if compiling from source):
- Python version:
- CUDA/cuDNN version (if applicable):
- GPU models and configuration (e.g. V100):
- Any other relevant information:

## Additional context

There is no code implementation for the rank=-1 case.

https://github.com/dmlc/dgl/blob/3d16000b4170fa741ed9e9667f22ba84d3493026/python/dgl/nn/pytorch/sparse_emb.py#L343-L347

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.