removes test task dir - NB it would be a good idea to add it in the f…

…uture

tests/tasks.py

@@ -503,3 +503,61 @@ def output_fn(input, target):
 
  layer_collection = LayerCollection.from_model(net)
  return (train_loader, layer_collection, net.parameters(), net, output_fn, 4)
+
+
+class LayerNormNet(nn.Module):
+ def __init__(self, out_size):
+ super(LayerNormNet, self).__init__()
+
+ self.linear1 = nn.Linear(18 * 18, out_size)
+ self.layer_norm1 = nn.LayerNorm((out_size,))
+
+ self.net = nn.Sequential(self.linear1, self.layer_norm1)
+
+ def forward(self, x):
+ x = x[:, :, 5:-5, 5:-5].contiguous()
+ x = x.view(x.size(0), -1)
+ return self.net(x)
+
+
+def get_layernorm_task():
+ train_set = get_mnist()
+ train_set = Subset(train_set, range(70))
+ train_loader = DataLoader(dataset=train_set, batch_size=30, shuffle=False)
+ net = LayerNormNet(out_size=3)
+ to_device_model(net)
+ net.eval()
+
+ def output_fn(input, target):
+ return net(to_device(input))
+
+ layer_collection = LayerCollection.from_model(net)
+ return (train_loader, layer_collection, net.parameters(), net, output_fn, 3)
+
+
+class LayerNormConvNet(nn.Module):
+ def __init__(self):
+ super(LayerNormConvNet, self).__init__()
+ self.layer = nn.Conv2d(1, 3, (3, 2), 2)
+ self.layer_norm = nn.LayerNorm((3, 8, 9))
+
+ def forward(self, x):
+ x = x[:, :, 5:-5, 5:-5]
+ x = self.layer(x)
+ x = self.layer_norm(x)
+ return x.sum(dim=(2, 3))
+
+
+def get_layernorm_conv_task():
+ train_set = get_mnist()
+ train_set = Subset(train_set, range(70))
+ train_loader = DataLoader(dataset=train_set, batch_size=30, shuffle=False)
+ net = LayerNormConvNet()
+ to_device_model(net)
+ net.eval()
+
+ def output_fn(input, target):
+ return net(to_device(input))
+
+ layer_collection = LayerCollection.from_model(net)
+ return (train_loader, layer_collection, net.parameters(), net, output_fn, 3)

tests/test_jacobian.py

@@ -12,12 +12,13 @@
  get_fullyconnect_onlylast_task,
  get_fullyconnect_task,
  get_fullyconnect_wn_task,
+ get_layernorm_conv_task,
+ get_layernorm_task,
  get_linear_conv_task,
  get_linear_fc_task,
  get_small_conv_transpose_task,
  get_small_conv_wn_task,
 )
-from test_tasks.layernorm import get_layernorm_conv_task, get_layernorm_task
 from utils import check_ratio, check_tensors
 
 from nngeometry.generator import Jacobian