TGN CUDA bugfix (#137)

* Bump MG version to 2.2.1 * Fix TGN to work with CUDA Co-authored-by: jmatak <[email protected]>
memgraph · Apr 20, 2022 · e6c024e · e6c024e
1 parent 6a7b8bd
commit e6c024e
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 17 deletions.
diff --git a/.github/workflows/docker_publish.yml b/.github/workflows/docker_publish.yml
@@ -12,7 +12,7 @@ jobs:
     env:
       DOCKER_ORGANIZATION_NAME: memgraph
       DOCKER_REPOSITORY_NAME: memgraph-mage
-      MEMGRAPH_VERSION: 2.2.0
+      MEMGRAPH_VERSION: 2.2.1
     steps:
       - name: Check out the repo
         uses: actions/checkout@v2
@@ -64,7 +64,7 @@ jobs:
       - name: Download memgraph binary
         run: |
           curl -L https://download.memgraph.com/memgraph/v${MEMGRAPH_VERSION}/debian-11/memgraph_${MEMGRAPH_VERSION}-1_amd64.deb > memgraph-amd64.deb
-          curl -L https://download.memgraph.com/memgraph/v${MEMGRAPH_VERSION}/debian-11-aarch64/memgraph_${MEMGRAPH_VERSION}-1_aarch64.deb > memgraph-arm64.deb
+          curl -L https://download.memgraph.com/memgraph/v${MEMGRAPH_VERSION}/debian-11-aarch64/memgraph_${MEMGRAPH_VERSION}-1_arm64.deb > memgraph-arm64.deb
 
       - name: Build & push prod docker images
         run: |

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -11,7 +11,7 @@ jobs:
   build:
     runs-on: ubuntu-latest
     env:
-      MEMGRAPH_VERSION: 2.2.0
+      MEMGRAPH_VERSION: 2.2.1
     strategy:
       matrix:
         architecture: ["amd64", "arm64"]
@@ -29,7 +29,7 @@ jobs:
       - name: Download memgraph binaries
         run: |
           curl -L https://download.memgraph.com/memgraph/v${MEMGRAPH_VERSION}/debian-11/memgraph_${MEMGRAPH_VERSION}-1_amd64.deb > memgraph-amd64.deb
-          curl -L https://download.memgraph.com/memgraph/v${MEMGRAPH_VERSION}/debian-11-aarch64/memgraph_${MEMGRAPH_VERSION}-1_aarch64.deb > memgraph-arm64.deb
+          curl -L https://download.memgraph.com/memgraph/v${MEMGRAPH_VERSION}/debian-11-aarch64/memgraph_${MEMGRAPH_VERSION}-1_arm64.deb > memgraph-arm64.deb
 
       - name: Build Memgraph MAGE:${{ matrix.target }}
         run: |

diff --git a/Dockerfile b/Dockerfile
@@ -4,7 +4,7 @@ FROM debian:bullseye as base
 
 USER root
 
-ARG MG_VERSION=2.2.0
+ARG MG_VERSION=2.2.1
 ARG PY_VERSION_DEFAULT
 ENV MG_VERSION ${MG_VERSION}
 ENV PY_VERSION ${PY_VERSION_DEFAULT}

diff --git a/python/tgn.py b/python/tgn.py
@@ -553,17 +553,18 @@ def process_batch_self_supervised() -> float:
         loss.backward()
         query_module_tgn.optimizer.step()
         query_module_tgn.m_loss.append(loss.item())
-
+    pos_prob_cpu = pos_prob.reshape((-1,)).detach().cpu()
+    neg_prob_cpu = neg_prob.reshape((-1,)).detach().cpu()
     # todo antoniofilipovic - update once we have logging API
     print(
         "POS PROB | NEG PROB",
-        pos_prob.reshape((-1,)).detach(),
-        neg_prob.reshape((-1,)).detach(),
+        pos_prob_cpu,
+        neg_prob_cpu,
     )
     pred_score = np.concatenate(
         [
-            (pos_prob.reshape((-1,))).detach().numpy(),
-            (neg_prob.reshape((-1,))).detach().numpy(),
+            pos_prob_cpu.numpy(),
+            neg_prob_cpu.numpy(),
         ]
     )
     true_label = np.concatenate(
@@ -577,10 +578,10 @@ def process_batch_self_supervised() -> float:
         np.sum(np.rint(true_label) == np.rint(pred_score)) * 1.0 / len(true_label)
     )
 
-    # update embeddings to newest ones that we can return on user request
+    # update embeddings to the newest ones that we can return on user request
     update_embeddings(
-        embeddings_source.cpu().detach().numpy(),
-        embeddings_dest.cpu().detach().numpy(),
+        embeddings_source.detach().cpu().numpy(),
+        embeddings_dest.detach().cpu().numpy(),
         sources,
         destinations,
     )
@@ -638,8 +639,8 @@ def process_batch_supervised() -> float:
 
     pred_score = np.concatenate(
         [
-            (src_prob.squeeze()).cpu().detach().numpy(),
-            (dest_prob.squeeze()).cpu().detach().numpy(),
+            (src_prob.squeeze()).detach().cpu().numpy(),
+            (dest_prob.squeeze()).detach().cpu().numpy(),
         ]
     )
     true_label = np.concatenate([np.array(labels[:, 0]), np.array(labels[:, 1])])
@@ -650,8 +651,8 @@ def process_batch_supervised() -> float:
 
     # update embeddings to newest ones that we can return on user request
     update_embeddings(
-        embeddings_source.cpu().detach().numpy(),
-        embeddings_dest.cpu().detach().numpy(),
+        embeddings_source.detach().cpu().numpy(),
+        embeddings_dest.cpu().detach().cpu().numpy(),
         sources,
         destinations,
     )