train_utils.py

import torch
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
import os, argparse, time
from scipy import sparse
from utils.losses import CELoss, MaskL2Loss, log_Bernoulli
from utils.util import ndcg_at_k
import random

import torch
import torch.nn as nn
import torch.nn.functional as F


def calc_distance(A, B, C):
	distances = (torch.sum(A**2, dim=1, keepdim=True) 
				+ torch.sum(B**2, dim=1)
				- 2 * torch.matmul(A, B.t()))
	distances = distances * C
	return distances.mean()

def h_calc_distance(A, B, C, split_dim=200):
	# Split A and B along the feature dimension (dim=1)
	distances = 0
	A1, A2 = torch.split(A, split_dim, dim=1)
	B1, B2 = torch.split(B, split_dim, dim=1)
	distances += calc_distance(A1, B1, C) 
	distances += calc_distance(A2, B2, C) 

	return distances

def run_epoch(model, embeddings, opts, device, load_data, samples_perc_per_epoch, batch_size, total_anneal_steps, anneal_cap, lambda_alignment, warmup_epoch, current_epoch, update_count, dropout_rate):
	

	model.train()
	embeddings.train()
	tmp_update_count = update_count

	for batch in generate(batch_size=batch_size, device=device, data_in=load_data, samples_perc_per_epoch=samples_perc_per_epoch, shuffle=True):
		X = batch.get_ratings_to_dev()


		logits, KLD, user_embedding = model(X, dropout_rate)
		logits_full, _, _ = model(X, 0.0)

		if model.model_type == 'hierachical_gated':
			reconstruction_loss = log_Bernoulli(X, logits, dim=1).mean() # specify for Hvamp
		else:
			reconstruction_loss = CELoss(X, logits)

		if model.prior_type == 'composition':
			beta = 0.005 # specify for RecVAE
		else:
			beta = min(anneal_cap, tmp_update_count / total_anneal_steps)

		if model.model_type == 'hierachical_gated':
			# pia = h_calc_distance
			pia = calc_distance
		else:
			pia = calc_distance

		loss = reconstruction_loss + beta * KLD

		alignment_loss = 0
		if lambda_alignment > 0.0:
			pairs = torch.nonzero(X).to(device)
			unique_item = pairs[:,1].unique()
			sample_item = torch.randint(0, len(unique_item), (X.shape[0],))
			personal_item = X[:,unique_item[sample_item]]
			item_embedding = embeddings(unique_item[sample_item])

			if current_epoch < warmup_epoch:
				# Since the items' embeddings are randomly initialized, it takes a few epochs to move them into the users' embeddings space.
				alignment_loss += lambda_alignment * pia(user_embedding.detach(), item_embedding, personal_item)
			else:
				alignment_loss += lambda_alignment * pia(user_embedding, item_embedding, personal_item)
			loss += alignment_loss	

		for optimizer in opts:
			optimizer.zero_grad()
		loss.backward()

		for optimizer in opts:
			optimizer.step()
		tmp_update_count += 1

	return tmp_update_count



def generate(batch_size, device, data_in, samples_perc_per_epoch, data_out=None, shuffle=False):
	assert 0 < samples_perc_per_epoch <= 1
	
	total_samples = data_in.shape[0]
	samples_per_epoch = int(total_samples * samples_perc_per_epoch)
	
	if shuffle:
		idxlist = np.arange(total_samples)
		np.random.shuffle(idxlist)
		idxlist = idxlist[:samples_per_epoch]
	else:
		idxlist = np.arange(samples_per_epoch)
	
	for st_idx in range(0, samples_per_epoch, batch_size):
		end_idx = min(st_idx + batch_size, samples_per_epoch)
		idx = idxlist[st_idx:end_idx]

		yield Batch(device, idx, data_in, data_out)


class Batch:
	def __init__(self, device, idx, data_in, data_out=None):
		self._device = device
		self._idx = idx
		self._data_in = data_in
		self._data_out = data_out
	
	def get_idx(self):
		return self._idx
	
	def get_idx_to_dev(self):
		return torch.LongTensor(self.get_idx()).to(self._device)
		
	def get_ratings(self, is_out=False):
		data = self._data_out if is_out else self._data_in
		return data[self._idx]
	
	def get_ratings_to_dev(self, is_out=False):
		return torch.Tensor(
			self.get_ratings(is_out).toarray()
		).to(self._device)

def evaluate(model, device, N, loaddata, batch_size, test_K=100):
	model.to(device)
	model.eval()
	# import pdb; pdb.set_trace()

	ndcg_list = []
	for st_idx in range(0, N, batch_size):
		end_idx = min(st_idx + batch_size, N)
		X_tr, X_te = loaddata(np.arange(st_idx, end_idx))
		
		## infer from partial observation X_tr
		with torch.no_grad():
			logits, _, _ = model(torch.FloatTensor(X_tr).to(device))

		logits = logits.cpu().numpy()
		logits[X_tr.nonzero()] = -np.inf
		## evaluate the inferred results with observation X_te
		ndcg_list.append(ndcg_at_k(logits, X_te, k=test_K))
	ndcg_list = np.concatenate(ndcg_list)
	ndcg = np.mean(ndcg_list)
	return ndcg