Sequential Recommendation for cold-start users with meta transitional learning
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

sampler.py 5.4KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143
  1. import sys
  2. import copy
  3. import torch
  4. import random
  5. import numpy as np
  6. from collections import defaultdict, Counter
  7. from multiprocessing import Process, Queue
  8. def random_neq(l, r, s, user_train,usernum):
  9. # t = np.random.randint(l, r)
  10. # while t in s:
  11. # t = np.random.randint(l, r)
  12. # return t
  13. user = np.random.randint(1, usernum + 1)
  14. candid_item = user_train[user][np.random.randint(0,len(user_train[user]))]
  15. while candid_item in s:
  16. user = np.random.randint(1, usernum + 1)
  17. candid_item = user_train[user][np.random.randint(0, len(user_train[user]))]
  18. return candid_item
  19. def sample_function_mixed(user_train, usernum, itemnum, batch_size, maxlen, result_queue, SEED):
  20. def sample():
  21. if random.random()<=1:
  22. user = np.random.randint(1, usernum + 1)
  23. while len(user_train[user]) <= 1: user = np.random.randint(1, usernum + 1)
  24. seq = np.zeros([maxlen], dtype=np.int32)
  25. pos = np.zeros([maxlen], dtype=np.int32)
  26. neg = np.zeros([maxlen*5], dtype=np.int32)
  27. if len(user_train[user]) < maxlen:
  28. nxt_idx = len(user_train[user]) - 1
  29. else:
  30. nxt_idx = np.random.randint(maxlen,len(user_train[user]))
  31. nxt = user_train[user][nxt_idx]
  32. idx = maxlen - 1
  33. ts = set(user_train[user])
  34. for i in reversed(user_train[user][min(0, nxt_idx - 1 - maxlen) : nxt_idx - 1]):
  35. seq[idx] = i
  36. pos[idx] = nxt
  37. # if nxt != 0: neg[idx] = random_neq(1, itemnum + 1, ts, user_train,usernum)
  38. nxt = i
  39. idx -= 1
  40. if idx == -1: break
  41. for i in range(len(neg)):
  42. neg[i] = random_neq(1, itemnum + 1, ts, user_train,usernum)
  43. curr_rel = user
  44. support_triples, support_negative_triples, query_triples, negative_triples = [], [], [], []
  45. for idx in range(maxlen-1):
  46. support_triples.append([seq[idx],curr_rel,pos[idx]])
  47. # support_negative_triples.append([seq[idx],curr_rel,neg[idx]])
  48. # support_negative_triples.append([seq[-1], curr_rel, neg[idx]])
  49. # for idx in range(maxlen*30 - 1):
  50. # support_negative_triples.append([seq[-1], curr_rel, neg[idx]])
  51. for j in range(5):
  52. for idx in range(maxlen-1):
  53. support_negative_triples.append([seq[idx], curr_rel, neg[j*maxlen + idx]])
  54. query_triples.append([seq[-1],curr_rel,pos[-1]])
  55. negative_triples.append([seq[-1],curr_rel,neg[-1]])
  56. return support_triples, support_negative_triples, query_triples, negative_triples, curr_rel
  57. else:
  58. user = np.random.randint(1, usernum + 1)
  59. while len(user_train[user]) <= 1: user = np.random.randint(1, usernum + 1)
  60. seq = np.zeros([maxlen], dtype=np.int32)
  61. pos = np.zeros([maxlen], dtype=np.int32)
  62. neg = np.zeros([maxlen], dtype=np.int32)
  63. list_idx = random.sample([i for i in range(len(user_train[user]))], maxlen + 1)
  64. list_item = [user_train[user][i] for i in sorted(list_idx)]
  65. nxt = list_item[-1]
  66. idx = maxlen - 1
  67. ts = set(user_train[user])
  68. for i in reversed(list_item[:-1]):
  69. seq[idx] = i
  70. pos[idx] = nxt
  71. if nxt != 0: neg[idx] = random_neq(1, itemnum + 1, ts)
  72. nxt = i
  73. idx -= 1
  74. if idx == -1: break
  75. curr_rel = user
  76. support_triples, support_negative_triples, query_triples, negative_triples = [], [], [], []
  77. for idx in range(maxlen-1):
  78. support_triples.append([seq[idx],curr_rel,pos[idx]])
  79. support_negative_triples.append([seq[idx],curr_rel,neg[idx]])
  80. query_triples.append([seq[-1],curr_rel,pos[-1]])
  81. negative_triples.append([seq[-1],curr_rel,neg[-1]])
  82. return support_triples, support_negative_triples, query_triples, negative_triples, curr_rel
  83. np.random.seed(SEED)
  84. while True:
  85. one_batch = []
  86. for i in range(batch_size):
  87. one_batch.append(sample())
  88. support, support_negative, query, negative, curr_rel = zip(*one_batch)
  89. result_queue.put(([support, support_negative, query, negative], curr_rel))
  90. class WarpSampler(object):
  91. def __init__(self, User, usernum, itemnum, batch_size=64, maxlen=10, n_workers=1):
  92. self.result_queue = Queue(maxsize=n_workers * 10)
  93. self.processors = []
  94. for i in range(n_workers):
  95. self.processors.append(
  96. Process(target=sample_function_mixed, args=(User,
  97. usernum,
  98. itemnum,
  99. batch_size,
  100. maxlen,
  101. self.result_queue,
  102. np.random.randint(2e9)
  103. )))
  104. self.processors[-1].daemon = True
  105. self.processors[-1].start()
  106. def next_batch(self):
  107. return self.result_queue.get()
  108. def close(self):
  109. for p in self.processors:
  110. p.terminate()
  111. p.join()