embedding work in progress

nngeometry/generator/jacobian/jacobian.py

@@ -312,14 +312,21 @@ def get_jacobian(self, examples):
  self.start = 0
  for d in loader:
  inputs = d[0]
- inputs.requires_grad = True
+ differentiate_wrt = []
+ if inputs.dtype in [
+ torch.float16,
+ torch.float32,
+ torch.float64,
+ ]:
+ inputs.requires_grad = True
+ differentiate_wrt.append(inputs)
  bs = inputs.size(0)
  output = self.function(*d).view(bs, self.n_output).sum(dim=0)
  for self.i_output in range(self.n_output):
  retain_graph = self.i_output < self.n_output - 1
  torch.autograd.grad(
  output[self.i_output],
- [inputs],
+ differentiate_wrt,
  retain_graph=retain_graph,
  only_inputs=True,
  )

nngeometry/layercollection.py

@@ -25,7 +25,8 @@ class LayerCollection:
  "Affine1d",
  "ConvTranspose2d",
  "Conv1d",
- "LayerNorm"
+ "LayerNorm",
+ "Embedding",
  ]
 
  def __init__(self, layers=None):
@@ -151,6 +152,10 @@ def _module_to_layer(mod):
  return LayerNormLayer(
  normalized_shape=mod.normalized_shape, bias=(mod.bias is not None)
  )
+ elif mod_class == "Embedding":
+ return EmbeddingLayer(
+ embedding_dim=mod.embedding_dim, num_embeddings=mod.num_embeddings
+ )
 
  def numel(self):
  """
@@ -292,6 +297,22 @@ def __eq__(self, other):
  )
 
 
+class EmbeddingLayer(AbstractLayer):
+ def __init__(self, num_embeddings, embedding_dim):
+ self.num_embeddings = num_embeddings
+ self.embedding_dim = embedding_dim
+ self.weight = Parameter(num_embeddings, embedding_dim)
+
+ def numel(self):
+ return self.weight.numel()
+
+ def __eq__(self, other):
+ return (
+ self.num_embeddings == other.num_embeddings
+ and self.embedding_dim == other.embedding_dim
+ )
+
+
 class BatchNorm1dLayer(AbstractLayer):
  def __init__(self, num_features):
  self.num_features = num_features

tests/tasks.py

@@ -4,7 +4,7 @@
 import torch.nn as nn
 import torch.nn.functional as tF
 from torch.nn.modules.conv import ConvTranspose2d
-from torch.utils.data import DataLoader, Subset
+from torch.utils.data import DataLoader, Subset, TensorDataset
 from torchvision import datasets, transforms
 
 from nngeometry.layercollection import LayerCollection
@@ -184,6 +184,33 @@ def output_fn(input, target):
  return (train_loader, layer_collection, net.parameters(), net, output_fn, 2)
 
 
+class EmbeddingNet(nn.Module):
+ def __init__(self):
+ super(EmbeddingNet, self).__init__()
+ self.embedding_layer = nn.Embedding(10, 3)
+
+ def forward(self, x):
+ output = self.embedding_layer(x)
+ print(output.size())
+ return output.sum(axis=1)
+
+
+def get_embedding_task():
+ train_set = TensorDataset(
+ torch.LongTensor([[1, 2, 4, 5], [4, 3, 2, 9]]), torch.LongTensor([2, 0])
+ )
+ train_loader = DataLoader(dataset=train_set, batch_size=2, shuffle=False)
+ net = EmbeddingNet()
+ to_device_model(net)
+ net.eval()
+
+ def output_fn(input, target):
+ return net(input)
+
+ layer_collection = LayerCollection.from_model(net)
+ return (train_loader, layer_collection, net.parameters(), net, output_fn, 3)
+
+
 class LinearConvNet(nn.Module):
  def __init__(self):
  super(LinearConvNet, self).__init__()

tests/test_jacobian.py

@@ -7,6 +7,7 @@
  get_conv_gn_task,
  get_conv_skip_task,
  get_conv_task,
+ get_embedding_task,
  get_fullyconnect_affine_task,
  get_fullyconnect_cosine_task,
  get_fullyconnect_onlylast_task,
@@ -35,6 +36,7 @@
 from nngeometry.object.vector import PVector, random_fvector, random_pvector
 
 linear_tasks = [
+ get_embedding_task,
  get_linear_fc_task,
  get_linear_conv_task,
  get_batchnorm_fc_linear_task,