From 72e663374b12b164dd544d6a980582e11f3affc4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire?= <gregmialon@hotmail.com>
Date: Mon, 10 Aug 2020 19:41:10 +0200
Subject: [PATCH 1/3] added QueryKernel, query_multihead_attn,
 QuerySeqAttention

---
 otk/layers.py   |  15 +++++-
 otk/models.py   | 139 +++++++++++++++++++++++++++++++++++++++++++++++-
 otk/sinkhorn.py |  25 +++++++++
 3 files changed, 177 insertions(+), 2 deletions(-)

diff --git a/otk/layers.py b/otk/layers.py
index 97cecc1..d65abb4 100644
--- a/otk/layers.py
+++ b/otk/layers.py
@@ -4,7 +4,7 @@
 from torch import nn
 import torch.optim as optim
 from .utils import spherical_kmeans, normalize
-from .sinkhorn import wasserstein_kmeans, multihead_attn
+from .sinkhorn import wasserstein_kmeans, multihead_attn, query_multihead_attn
 
 
 class OTKernel(nn.Module):
@@ -165,3 +165,16 @@ def score(self, X, y):
             scores = scores.argmax(-1)
             scores = scores.cpu()
         return torch.mean((scores == y).float()).item()
+
+
+class QueryKernel(OTKernel):
+    def get_attn(self, input, mask=None, position_filter=None):
+        """Compute the attention weight using Sinkhorn OT
+        input: batch_size x in_size x in_dim
+        mask: batch_size x in_size
+        self.weight: heads x out_size x in_dim
+        output: batch_size x (out_size x heads) x in_size
+        """
+        return query_multihead_attn(
+            input, self.weight, mask=mask,
+            position_filter=position_filter)
\ No newline at end of file
diff --git a/otk/models.py b/otk/models.py
index 1188c10..e52dd46 100644
--- a/otk/models.py
+++ b/otk/models.py
@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import torch
 from torch import nn
-from .layers import OTKernel, Linear
+from .layers import OTKernel, Linear, QueryKernel
 from ckn.layers import BioEmbedding
 from ckn.models import CKNSequential
 
@@ -142,3 +142,140 @@ def unsup_train(self, data_loader, n_sampling_patches=300000, n_samples=5000, wb
             cur_samples += size
         print(patches.shape)
         self.attention.unsup_train(patches, wb=wb, use_cuda=use_cuda)
+
+class QuerySeqAttention(nn.Module):
+    def __init__(self, in_channels, nclass, hidden_sizes, filter_sizes,
+                 subsamplings, kernel_args=None, eps=0.1, heads=1,
+                 out_size=1, max_iter=50, alpha=0., fit_bias=True, mask_zeros=True):
+        super().__init__()
+        self.embed_layer = BioEmbedding(
+            in_channels, False, mask_zeros=True, no_embed=True)
+        self.ckn_model = CKNSequential(
+            in_channels, hidden_sizes, filter_sizes,
+            subsamplings, kernel_args_list=kernel_args)
+        self.attention = QueryKernel(hidden_sizes[-1], out_size, heads=heads,
+                                     eps=eps, max_iter=max_iter)
+        self.out_features = out_size * heads * hidden_sizes[-1]
+        self.nclass = nclass
+
+        self.classifier = Linear(self.out_features, nclass, bias=fit_bias)
+        self.alpha = alpha
+        self.mask_zeros = mask_zeros
+
+    def feature_parameters(self):
+        import itertools
+        return itertools.chain(self.ckn_model.parameters(), self.attention.parameters())
+
+    def normalize_(self):
+        self.ckn_model.normalize_()
+
+    def ckn_representation_at(self, input, n=0):
+        output = self.embed_layer(input)
+        mask = self.embed_layer.compute_mask(input)
+        output = self.ckn_model.representation(output, n)
+        mask = self.ckn_model.compute_mask(mask, n)
+        return output, mask
+
+    def ckn_representation(self, input):
+        output = self.embed_layer(input)
+        output = self.ckn_model(output).permute(0, 2, 1).contiguous()
+        return output
+
+    def representation(self, input):
+        output = self.embed_layer(input)
+        mask = self.embed_layer.compute_mask(input)
+        output = self.ckn_model(output).permute(0, 2, 1).contiguous()
+        mask = self.ckn_model.compute_mask(mask)
+        if not self.mask_zeros:
+            mask = None
+        output = self.attention(output, mask).reshape(output.shape[0], -1)
+        return output
+
+    def forward(self, input):
+        output = self.representation(input)
+        return self.classifier(output)
+
+    def predict(self, data_loader, only_repr=False, use_cuda=False):
+        n_samples = len(data_loader.dataset)
+        target_output = torch.LongTensor(n_samples)
+        batch_start = 0
+        for i, (data, target) in enumerate(data_loader):
+            batch_size = data.shape[0]
+            if use_cuda:
+                data = data.cuda()
+            with torch.no_grad():
+                if only_repr:
+                    batch_out = self.representation(data).data.cpu()
+                else:
+                    batch_out = self(data).data.cpu()
+            if i == 0:
+                output = batch_out.new_empty([n_samples] + list(batch_out.shape[1:]))
+            output[batch_start:batch_start + batch_size] = batch_out
+            target_output[batch_start:batch_start + batch_size] = target
+            batch_start += batch_size
+        return output, target_output
+
+    def train_classifier(self, data_loader, criterion=None, epochs=100, optimizer=None, use_cuda=False):
+        encoded_train, encoded_target = self.predict(
+            data_loader, only_repr=True, use_cuda=use_cuda)
+        self.classifier.fit(encoded_train, encoded_target, criterion,
+                            reg=self.alpha, epochs=epochs, optimizer=optimizer, use_cuda=use_cuda)
+
+    def unsup_train(self, data_loader, n_sampling_patches=300000, n_samples=5000, wb=False, use_cuda=False):
+        self.eval()
+        if use_cuda:
+            self.cuda()
+
+        for i, ckn_layer in enumerate(self.ckn_model):
+            print("Training ckn layer {}".format(i))
+            n_patches = 0
+            try:
+                n_patches_per_batch = (n_sampling_patches + len(data_loader) - 1) // len(data_loader)
+            except:
+                n_patches_per_batch = 1000
+            patches = torch.Tensor(n_sampling_patches, ckn_layer.patch_dim)
+            if use_cuda:
+                patches = patches.cuda()
+
+            for data, _ in data_loader:
+                if n_patches >= n_sampling_patches:
+                    continue
+                if use_cuda:
+                    data = data.cuda()
+                with torch.no_grad():
+                    data, mask = self.ckn_representation_at(data, i)
+                    data_patches = ckn_layer.sample_patches(
+                        data, mask, n_patches_per_batch)
+                size = data_patches.size(0)
+                if n_patches + size > n_sampling_patches:
+                    size = n_sampling_patches - n_patches
+                    data_patches = data_patches[:size]
+                patches[n_patches: n_patches + size] = data_patches
+                n_patches += size
+
+            print("total number of patches: {}".format(n_patches))
+            patches = patches[:n_patches]
+            ckn_layer.unsup_train(patches, init=None)
+
+        n_samples = min(n_samples, len(data_loader.dataset))
+        cur_samples = 0
+        print("Training attention layer")
+        for i, (data, _) in enumerate(data_loader):
+            if cur_samples >= n_samples:
+                continue
+            if use_cuda:
+                data = data.cuda()
+            with torch.no_grad():
+                data = self.ckn_representation(data)
+
+            if i == 0:
+                patches = torch.empty([n_samples]+list(data.shape[1:]))
+
+            size = data.shape[0]
+            if cur_samples + size > n_samples:
+                size = n_samples - cur_samples
+                data = data[:size]
+            patches[cur_samples: cur_samples + size] = data
+            cur_samples += size
+        print(patches.shape)
+        self.attention.unsup_train(patches, wb=wb, use_cuda=use_cuda)
diff --git a/otk/sinkhorn.py b/otk/sinkhorn.py
index 581f7ae..48c1161 100644
--- a/otk/sinkhorn.py
+++ b/otk/sinkhorn.py
@@ -98,6 +98,31 @@ def multihead_attn(input, weight, mask=None, eps=1.0, return_kernel=False,
     K = K.permute(0, 3, 1, 2).contiguous()
     return K
 
+def query_multihead_attn(input, weight, mask=None, return_kernel=False,
+                   position_filter=None):
+    """Comput the attention weight using the references as queries
+    input: n x in_size x in_dim
+    mask: n x in_size
+    weight: m x out_size x in_dim (m: number of heads/ref)
+    output: n x out_size x m x in_size
+    """
+    n, in_size, in_dim = input.shape
+    m, out_size = weight.shape[:-1]
+    K = torch.tensordot(input, weight, dims=[[-1], [-1]])
+    K = K.permute(0, 2, 1, 3)
+    if position_filter is not None:
+        K = position_filter * K
+    # K: n x m x in_size x out_size
+    K = K.reshape(-1, in_size, out_size)
+    # K: nm x in_size x out_size
+    if return_kernel:
+        return K.reshape(n, m)
+    K = K.reshape(n, m, in_size, out_size)
+    if position_filter is not None:
+        K = position_filter * K
+    K = K.permute(0, 3, 1, 2).contiguous()
+    return K
+
 def wasserstein_barycenter(x, c, eps=1.0, max_iter=100, sinkhorn_iter=50, log_domain=False):
     """
     x: n x in_size x in_dim

From c27b7e6587fdeb0a131378f62f39b2b7d7d26a65 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire?= <gregmialon@hotmail.com>
Date: Mon, 10 Aug 2020 20:35:07 +0200
Subject: [PATCH 2/3] added query baseline

---
 experiments/scop175_sup.py | 21 +++++++++++++++------
 1 file changed, 15 insertions(+), 6 deletions(-)

diff --git a/experiments/scop175_sup.py b/experiments/scop175_sup.py
index e8a076b..034bd81 100644
--- a/experiments/scop175_sup.py
+++ b/experiments/scop175_sup.py
@@ -2,7 +2,7 @@
 import argparse
 
 from ckn.data.loader_scop import load_data
-from otk.models import SeqAttention
+from otk.models import SeqAttention, QuerySeqAttention
 from torch.utils.data import DataLoader
 import torch
 from torch import nn
@@ -67,6 +67,8 @@ def load_args():
         help='initial learning rate')
     parser.add_argument("--alternating", action='store_true',
         help='alternating training')
+    parser.add_argument("--baseline", type=str, default='ours',
+                        choices=['ours', 'query'])
     args = parser.parse_args()
     args.use_cuda = torch.cuda.is_available()
     # check shape
@@ -232,11 +234,18 @@ def main():
         val_loader = DataLoader(
             val_dset, batch_size=args.batch_size, shuffle=False, **loader_args)
 
-    model = SeqAttention(
-        45, 1195, args.n_filters, args.len_motifs, args.subsamplings,
-        kernel_args=args.kernel_params, alpha=args.weight_decay,
-        eps=args.eps, heads=args.heads, out_size=args.out_size,
-        max_iter=args.max_iter)
+    if args.baseline == 'query':
+        model = QuerySeqAttention(
+            45, 1195, args.n_filters, args.len_motifs, args.subsamplings,
+            kernel_args=args.kernel_params, alpha=args.weight_decay,
+            eps=args.eps, heads=args.heads, out_size=args.out_size,
+            max_iter=args.max_iter)
+    else:
+        model = SeqAttention(
+            45, 1195, args.n_filters, args.len_motifs, args.subsamplings,
+            kernel_args=args.kernel_params, alpha=args.weight_decay,
+            eps=args.eps, heads=args.heads, out_size=args.out_size,
+            max_iter=args.max_iter)
     print(model)
     print(len(train_dset))
 

From 0974996c43cc24b26cadcc63d6c7b96ede33042e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire?= <gregmialon@hotmail.com>
Date: Mon, 10 Aug 2020 23:30:51 +0200
Subject: [PATCH 3/3] query att works

---
 otk/sinkhorn.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/otk/sinkhorn.py b/otk/sinkhorn.py
index 48c1161..b01027f 100644
--- a/otk/sinkhorn.py
+++ b/otk/sinkhorn.py
@@ -120,8 +120,8 @@ def query_multihead_attn(input, weight, mask=None, return_kernel=False,
     K = K.reshape(n, m, in_size, out_size)
     if position_filter is not None:
         K = position_filter * K
-    K = K.permute(0, 3, 1, 2).contiguous()
-    return K
+    K = torch.nn.Softmax(dim=2)(K / math.sqrt(in_dim))
+    return K.permute(0, 3, 1, 2).contiguous()
 
 def wasserstein_barycenter(x, c, eps=1.0, max_iter=100, sinkhorn_iter=50, log_domain=False):
     """