nn.Embedding weight normalization error in CUDA?

Sam-gege: <p>Hi, it seems that weight renormalized in cpu and gpu gives different result:</p>
<pre><code class="lang-python">import torch
import numpy as np

np.random.seed(1)

# initialize input and embedding weight
max_norm = 1
x_np = np.random.randint(0, 1000, (128, 20))
weight_np = np.random.randn(1000, 100).astype(np.float32)

# cpu
m_cpu = torch.nn.Embedding(1000, 100, max_norm=max_norm)
m_cpu.weight.data = torch.tensor(weight_np)
y_cpu = m_cpu(torch.tensor(x_np))

# gpu
m_cuda = torch.nn.Embedding(1000, 100, max_norm=max_norm).cuda()
m_cuda.weight.data = torch.tensor(weight_np).cuda()
y_cuda = m_cuda(torch.tensor(x_np).cuda())

# normalize weight manually
indices = np.unique(x_np)
norm = np.linalg.norm(weight_np[indices], ord=2, axis=-1, keepdims=True)  # ord = 2 is the default of nn.Embedding
rescale = (norm &gt; max_norm)[:, 0]
selection = np.zeros(weight_np.shape[0], dtype=bool)
selection[indices] = rescale
weight_np[selection] *= max_norm / (norm[rescale] + 1e-7)

# max diff
print(torch.abs(m_cpu.weight - m_cuda.weight.cpu()).max())  # max diff between weight renormalized in gpu and cpu
print(torch.abs(m_cpu.weight - torch.tensor(weight_np)).max())  # weight renormalized in cpu is correct
</code></pre>
<p>The output is:</p>
<pre><code class="lang-python">tensor(0.3130, grad_fn=&lt;MaxBackward1&gt;)
tensor(8.9407e-08, grad_fn=&lt;MaxBackward1&gt;)
</code></pre>
<p>Am I doing something wrong? Thanks!</p>

ptrblck: <p>The manipulation of the <code>.data</code> attribute is deprecated and can yield unwanted side effects.<br>
Copy the values into the <code>weigth</code> parameter and load the <code>state_dict</code> instead:</p>
<pre><code class="lang-python">import torch
import numpy as np

np.random.seed(1)

# initialize input and embedding weight
max_norm = 1
x_np = np.random.randint(0, 1000, (128, 20))
weight_np = np.random.randn(1000, 100).astype(np.float32)

# cpu
m_cpu = torch.nn.Embedding(1000, 100, max_norm=max_norm)
with torch.no_grad():
    m_cpu.weight.copy_(torch.tensor(weight_np))
y_cpu = m_cpu(torch.tensor(x_np))

# gpu
m_cuda = torch.nn.Embedding(1000, 100, max_norm=max_norm)#.cuda()
m_cuda.load_state_dict(m_cpu.state_dict())
m_cuda.cuda()
y_cuda = m_cuda(torch.tensor(x_np).cuda())

# normalize weight manually
indices = np.unique(x_np)
norm = np.linalg.norm(weight_np[indices], ord=2, axis=-1, keepdims=True)  # ord = 2 is the default of nn.Embedding
rescale = (norm &gt; max_norm)[:, 0]
selection = np.zeros(weight_np.shape[0], dtype=bool)
selection[indices] = rescale
weight_np[selection] *= max_norm / (norm[rescale] + 1e-7)

# max diff
print(torch.abs(m_cpu.weight - m_cuda.weight.cpu()).max())  # max diff between weight renormalized in gpu and cpu
print(torch.abs(m_cpu.weight - torch.tensor(weight_np)).max())  
</code></pre>
<p>Output:</p>
<pre><code class="lang-python">tensor(2.3842e-07, grad_fn=&lt;MaxBackward1&gt;)
tensor(8.9407e-08, grad_fn=&lt;MaxBackward1&gt;)
</code></pre>

Sam-gege: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>, suppose I only have one instance of <code>nn.Embedding</code>, i.e. no load <code>state_dict</code> from other instance:</p>
<pre><code class="lang-python">import torch
import numpy as np

np.random.seed(1)

# initialize input and embedding weight
max_norm = 1
x_np = np.random.randint(0, 1000, (128, 20))
weight_np = np.random.randn(1000, 100).astype(np.float32)

# gpu
m = torch.nn.Embedding(1000, 100, max_norm=max_norm)
## replace weight
????????????????????
##
m_cuda = m.cuda()
y_cuda = m_cuda(torch.tensor(x_np).cuda())


# correct normalize weight
indices = np.unique(x_np)
norm = np.linalg.norm(weight_np[indices], ord=2, axis=-1, keepdims=True)
rescale = (norm &gt; max_norm)[:, 0]
selection = np.zeros(weight_np.shape[0], dtype=bool)
selection[indices] = rescale
weight_np[selection] *= max_norm / (norm[rescale] + 1e-7)

# max diff
print(torch.abs(torch.tensor(weight_np) - m_cuda.weight.cpu()).max())
</code></pre>
<p>I’ve tried to replace weight with</p>
<pre><code class="lang-python">with torch.no_grad():
    m.weight.copy_(torch.tensor(weight_np))
</code></pre>
<p>but the result is wrong as before.<br>
Are you suggesting that, to replace the weight, I need to create another <code>nn.Embedding</code> instance, replace its weight, then load its <code>state_dict</code>?</p>
<p>Thanks for your help!</p>