class_docstr: Knowledge graph embedding application.
  
  Given a knowledge graph, it embeds each entity and relation into a continuous vector representation respectively.
  The learned embeddings can be used for analysis of knowledge graphs.
  e.g. **entity prediction**, **link prediction**.
  The likelihood of edges can be predicted by computing the score function over embeddings of triplets.
  
  Supported Models:
      - TransE (`Translating Embeddings for Modeling Multi-relational Data`_)
      - DistMult (`Embedding Entities and Relations for Learning and Inference in Knowledge Bases`_)
      - ComplEx (`Complex Embeddings for Simple Link Prediction`_)
      - SimplE (`SimplE Embedding for Link Prediction in Knowledge Graphs`_)
      - RotatE (`RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space`_)
      - QuatE (`Quaternion Knowledge Graph Embeddings`)
  
  .. _Translating Embeddings for Modeling Multi-relational Data:
      http://papers.nips.cc/paper/5071-translating-embeddings-for-modeling-multi-relational-data.pdf
  .. _Embedding Entities and Relations for Learning and Inference in Knowledge Bases:
      https://arxiv.org/pdf/1412.6575.pdf
  .. _Complex Embeddings for Simple Link Prediction:
      http://proceedings.mlr.press/v48/trouillon16.pdf
  .. _SimplE Embedding for Link Prediction in Knowledge Graphs:
      https://papers.nips.cc/paper/7682-simple-embedding-for-link-prediction-in-knowledge-graphs.pdf
  .. _RotatE\: Knowledge Graph Embedding by Relational Rotation in Complex Space:
      https://arxiv.org/pdf/1902.10197.pdf
  .. _Quaternion Knowledge Graph Embeddings:
      https://papers.nips.cc/paper/8541-quaternion-knowledge-graph-embeddings.pdf
  
  Parameters:
      dim (int): dimension of embeddings
      gpus (list of int, optional): GPU ids, default is all GPUs
      cpu_per_gpu (int, optional): number of CPU threads per GPU, default is all CPUs
      float_type (dtype, optional): type of parameters
      index_type (dtype, optional): type of graph indexes
  
  Note:
      The implementation of TransE, DistMult, ComplEx, SimplE and QuatE are slightly different from their original
      papers.
      The loss function and the regularization term generally follow `this repo`_.
      Self-adversarial negative sampling is also adopted in these models like RotatE.
  
  .. _this repo: https://github.com/DeepGraphLearning/KnowledgeGraphEmbedding
  
  See also:
      :class:`KnowledgeGraph <graphvite.graph.KnowledgeGraph>`,
      :class:`KnowledgeGraphSolver <graphvite.solver.KnowledgeGraphSolver>`
class_docstr_tok: ["Knowledge","graph","embedding","application",".","Given","a","knowledge","graph",",","it","embeds","each","entity","and","relation","into","a","continuous","vector","representation","respectively",".","The","learned","embeddings","can","be","used","for","analysis","of","knowledge","graphs",".","e.g",".","*","*","entity","prediction","*","*",",","*","*","link","prediction","*","*",".","The","likelihood","of","edges","can","be","predicted","by","computing","the","score","function","over","embeddings","of","triplets",".","Supported","Models",":","-","TransE","(","`","Translating","Embeddings","for","Modeling","Multi-relational","Data","`","_",")","-","DistMult","(","`","Embedding","Entities","and","Relations","for","Learning","and","Inference","in","Knowledge","Bases","`","_",")","-","ComplEx","(","`","Complex","Embeddings","for","Simple","Link","Prediction","`","_",")","-","SimplE","(","`","SimplE","Embedding","for","Link","Prediction","in","Knowledge","Graphs","`","_",")","-","RotatE","(","`","RotatE",":","Knowledge","Graph","Embedding","by","Relational","Rotation","in","Complex","Space","`","_",")","-","QuatE","(","`","Quaternion","Knowledge","Graph","Embeddings","`",")","..","_Translating","Embeddings","for","Modeling","Multi-relational","Data",":","http",":","\/\/papers.nips.cc\/paper\/5071-translating-embeddings-for-modeling-multi-relational-data.pdf","..","_Embedding","Entities","and","Relations","for","Learning","and","Inference","in","Knowledge","Bases",":","https",":","\/\/arxiv.org\/pdf\/1412.6575.pdf","..","_Complex","Embeddings","for","Simple","Link","Prediction",":","http",":","\/\/proceedings.mlr.press\/v48\/trouillon16.pdf","..","_SimplE","Embedding","for","Link","Prediction","in","Knowledge","Graphs",":","https",":","\/\/papers.nips.cc\/paper\/7682-simple-embedding-for-link-prediction-in-knowledge-graphs.pdf","..","_RotatE\\",":","Knowledge","Graph","Embedding","by","Relational","Rotation","in","Complex","Space",":","https",":","\/\/arxiv.org\/pdf\/1902.10197.pdf","..","_Quaternion","Knowledge","Graph","Embeddings",":","https",":","\/\/papers.nips.cc\/paper\/8541-quaternion-knowledge-graph-embeddings.pdf","Parameters",":","dim","(","int",")",":","dimension","of","embeddings","gpus","(","list","of","int",",","optional",")",":","GPU","ids",",","default","is","all","GPUs","cpu_per_gpu","(","int",",","optional",")",":","number","of","CPU","threads","per","GPU",",","default","is","all","CPUs","float_type","(","dtype",",","optional",")",":","type","of","parameters","index_type","(","dtype",",","optional",")",":","type","of","graph","indexes","Note",":","The","implementation","of","TransE",",","DistMult",",","ComplEx",",","SimplE","and","QuatE","are","slightly","different","from","their","original","papers",".","The","loss","function","and","the","regularization","term","generally","follow","`","this","repo","`","_",".","Self-adversarial","negative","sampling","is","also","adopted","in","these","models","like","RotatE",".","..","_this","repo",":","https",":","\/\/github.com\/DeepGraphLearning\/KnowledgeGraphEmbedding","See","also",":",":","class",":","`","KnowledgeGraph","<","graphvite.graph.KnowledgeGraph",">","`",",",":","class",":","`","KnowledgeGraphSolver","<","graphvite.solver.KnowledgeGraphSolver",">","`"]
code: def get_batch_size(self, sample_size):
      import psutil
  
      memory = psutil.virtual_memory()
  
      batch_size = int(
          self.SAMPLE_PER_DIMENSION
          * self.dim
          * self.graph.num_vertex
          * self.solver.num_partition
          / self.solver.num_worker
          / sample_size
      )
      # 2 triplet (Python, C++ sample pool) + 1 sample index
      mem_per_sample = sample_size * (
          2 * 3 * np.uint32().itemsize + 1 * np.uint64().itemsize
      )
      max_batch_size = int(
          memory.available / mem_per_sample / self.MEMORY_SCALE_FACTOR
      )
      if max_batch_size < batch_size:
          logger.info(
              "Memory is not enough for optimal prediction batch size. "
              "Use the maximal possible size instead."
          )
          batch_size = max_batch_size
      return batch_size
  
imports: from __future__ import print_function, absolute_import, unicode_literals, division
  import os
  import re
  import pickle
  import logging
  import multiprocessing
  from collections import defaultdict
  from future.builtins import str, map, range
  from easydict import EasyDict
  import numpy
  from ..None import lib, cfg, auto
  from ..None import graph, solver
  from ..util import assert_in, monitor, SharedNDArray
label_desc: Use image node_id 6 for calling the KnowledgeGraphApplication obj's underlying member method code with example usage: obj.get_batch_size(sample_size) and returns: batch_size