def MTTKRP_TTTP(self, sk, out):
if self.use_MTTKRP:
if self.string == "U":
ctf.MTTKRP(ctf.TTTP(self.omega, [sk, self.f1, self.f2]),
[out, self.f1, self.f2], 0)
elif self.string == "V":
ctf.MTTKRP(ctf.TTTP(self.omega, [self.f1, sk, self.f2]),
[self.f1, out, self.f2], 1)
elif self.string == "W":
ctf.MTTKRP(ctf.TTTP(self.omega, [self.f1, self.f2, sk]),
[self.f1, self.f2, out], 2)
else:
print("Invalid string for implicit MTTKRP_TTTP")
else:
idx = "ir"
if self.string == "U":
out.i(idx) << self.f1.i("J"+idx[1]) \
*self.f2.i("K"+idx[1]) \
*ctf.TTTP(self.omega, [sk, self.f1, self.f2]).i(idx[0]+"JK")
if self.string == "V":
out.i(idx) << self.f1.i("I"+idx[1]) \
*self.f2.i("K"+idx[1]) \
*ctf.TTTP(self.omega, [self.f1, sk, self.f2]).i("I"+idx[0]+"K")
if self.string == "W":
out.i(idx) << self.f1.i("I"+idx[1]) \
*self.f2.i("J"+idx[1]) \
*ctf.TTTP(self.omega, [self.f1, self.f2, sk]).i("IJ"+idx[0])
def sparse_update(T, factors, Lambda, sizes, rank, stepSize, sample_rate, times, use_MTTKRP):
starting_time = time.time()
t_go = ctf.timer("SGD_getOmega")
t_go.start()
omega = getOmega(T)
t_go.stop()
dimension = len(sizes)
indexes = INDEX_STRING[:dimension]
R = ctf.tensor(copy=T) #ctf.tensor(tuple(sizes), sp = True)
times[2] += time.time() - starting_time
for i in range(dimension):
starting_time = time.time()
R.i(indexes) << -1.* ctf.TTTP(omega, factors).i(indexes)
times[3] += time.time() - starting_time
starting_time = time.time()
times[4] += time.time() - starting_time
starting_time = time.time()
times[5] += time.time() - starting_time
starting_time = time.time()
t_ctr = ctf.timer("SGD_main_contraction")
t_ctr.start()
if use_MTTKRP:
new_fi = (1- stepSize * 2 * Lambda * sample_rate)*factors[i]
ctf.MTTKRP(R, factors, i)
stepSize*factors[i].i("ir") << new_fi.i("ir")
else:
tup_list = [factors[i].i(indexes[i] + "r") for i in range(dimension)]
Hterm = reduce(lambda x, y: x * y, tup_list[:i] + tup_list[i + 1:])
(1- stepSize * 2 * Lambda * sample_rate)*factors[i].i(indexes[i] + "r") << stepSize * Hterm * R.i(indexes)
t_ctr.stop()
times[6] += time.time() - starting_time
if i < dimension - 1:
R = ctf.tensor(copy=T)
def test_MTTKRP_vec(self):
for N in range(2, 5):
lens = numpy.random.randint(3, 4, N)
A = ctf.tensor(lens)
A.fill_sp_random(-1., 1., .5)
mats = []
for i in range(N):
mats.append(ctf.random.random([lens[i]]))
for i in range(N):
ctr = A.i("ijklm"[0:N])
for j in range(N):
if i != j:
ctr *= mats[j].i("ijklm"[j])
ans = ctf.zeros(mats[i].shape)
ans.i("ijklm"[i]) << ctr
ctf.MTTKRP(A, mats, i)
self.assertTrue(allclose(ans, mats[i]))
def MTTKRP(T, A, idx):
return ctf.MTTKRP(T, A, idx)
def run_bench(num_iter, s_start, s_end, mult, R, sp, sp_init, use_cust_MTTKRP):
wrld = ctf.comm()
s = s_start
nnz = float(s_start*s_start*s_start)*sp_init
agg_s = []
agg_avg_times = []
agg_min_times = []
agg_max_times = []
agg_min_95 = []
agg_max_95 = []
if num_iter > 1:
if ctf.comm().rank() == 0:
print("Performing MTTKRP WARMUP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
T = ctf.tensor((s,s,s),sp=sp)
T.fill_sp_random(-1.,1.,float(nnz)/float(s*s*s))
U = ctf.random.random((s,R))
V = ctf.random.random((s,R))
W = ctf.random.random((s,R))
if use_cust_MTTKRP:
ctf.MTTKRP(T,[U,V,W],0)
ctf.MTTKRP(T,[U,V,W],1)
ctf.MTTKRP(T,[U,V,W],2)
else:
U = ctf.einsum("ijk,jr,kr->ir",T,V,W)
V = ctf.einsum("ijk,ir,kr->jr",T,U,W)
W = ctf.einsum("ijk,ir,jr->kr",T,U,V)
if ctf.comm().rank() == 0:
print("Completed MTTKRP WARMUP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
while s<=s_end:
agg_s.append(s)
T = ctf.tensor((s,s,s),sp=sp)
T.fill_sp_random(-1.,1.,float(nnz)/float(s*s*s))
if ctf.comm().rank() == 0:
print("Performing MTTKRP with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
U = ctf.random.random((s,R))
V = ctf.random.random((s,R))
W = ctf.random.random((s,R))
te1 = 0.
te2 = 0.
te3 = 0.
avg_times = []
for i in range(num_iter):
t0 = time.time()
if use_cust_MTTKRP == 1:
ctf.MTTKRP(T,[U,V,W],0)
else:
U = ctf.einsum("ijk,jr,kr->ir",T,V,W)
t1 = time.time()
ite1 = t1 - t0
te1 += ite1
t0 = time.time()
if use_cust_MTTKRP == 1:
ctf.MTTKRP(T,[U,V,W],1)
else:
V = ctf.einsum("ijk,ir,kr->jr",T,U,W)
t1 = time.time()
ite2 = t1 - t0
te2 += ite2
t0 = time.time()
if use_cust_MTTKRP == 1:
ctf.MTTKRP(T,[U,V,W],2)
else:
W = ctf.einsum("ijk,ir,jr->kr",T,U,V)
t1 = time.time()
ite3 = t1 - t0
te3 += ite3
if ctf.comm().rank() == 0:
print(ite1,ite2,ite3,"avg:",(ite1+ite2+ite3)/3.)
avg_times.append((ite1+ite2+ite3)/3.)
if ctf.comm().rank() == 0:
print("Completed",num_iter,"iterations, took",te1/num_iter,te2/num_iter,te3/num_iter,"seconds on average for 3 variants.")
avg_time = (te1+te2+te3)/(3*num_iter)
agg_avg_times.append(avg_time)
print("MTTKRP took",avg_times,"seconds on average across variants with s =",s,"nnz =",nnz,"sp =",sp,"sp_init =",sp_init,"use_cust_MTTKRP=",use_cust_MTTKRP)
min_time = np.min(avg_times)
max_time = np.max(avg_times)
agg_min_times.append(min_time)
agg_max_times.append(max_time)
print("min/max interval is [",min_time,",",max_time,"]")
stddev = np.std(avg_times)
min_95 = (te1+te2+te3)/(3*num_iter)-2*stddev
max_95 = (te1+te2+te3)/(3*num_iter)+2*stddev
agg_min_95.append(min_95)
agg_max_95.append(max_95)
print("95% confidence interval is [",min_95,",",max_95,"]")
s = int(s*mult)
if ctf.comm().rank() == 0:
print("s min_time min_95 avg_time max_95 max_time")
for i in range(len(agg_s)):
print(agg_s[i], agg_min_times[i], agg_min_95[i], agg_avg_times[i], agg_max_95[i], agg_max_times[i])
def updateFactor(T, U, V, W, regParam, omega, I, J, K, r, block_size, string,
use_implicit, use_MTTKRP):
t_RHS = ctf.timer("ALS_imp_cg_RHS")
t_cg_TTTP = ctf.timer("ALS_imp_cg_TTTP")
t_o_slice = ctf.timer("ALS_imp_omega_slice")
t_form_EQs = ctf.timer("ALS_exp_form_EQs")
t_form_RHS = ctf.timer("ALS_exp_form_RHS")
if (string == "U"):
num_blocks = int((I + block_size - 1) / block_size)
for n in range(num_blocks):
I_start = n * block_size
I_end = min(I, I_start + block_size)
bsize = I_end - I_start
t_o_slice.start()
if num_blocks > 1:
nomega = omega[I_start:I_end, :, :]
else:
nomega = omega
t_o_slice.stop()
x0 = ctf.random.random((bsize, r))
b = ctf.tensor((bsize, r))
t_RHS.start()
if num_blocks == 1:
if use_MTTKRP:
ctf.MTTKRP(T, [b, V, W], 0)
else:
b.i("ir") << V.i("Jr") * W.i("Kr") * T.i("iJK")
else:
if use_MTTKRP:
ctf.MTTKRP(T[I_start:I_end, :, :], [b, V, W], 0)
else:
b.i("ir") << V.i("Jr") * W.i("Kr") * T[
I_start:I_end, :, :].i("iJK") # RHS; ATb
t_RHS.stop()
if use_implicit:
Ax0 = ctf.tensor((bsize, r))
t_cg_TTTP.start()
if use_MTTKRP:
ctf.MTTKRP(ctf.TTTP(nomega, [x0, V, W]), [Ax0, V, W], 0)
else:
Ax0.i("ir") << V.i("Jr") * W.i("Kr") * ctf.TTTP(
nomega, [x0, V, W]).i("iJK")
t_cg_TTTP.stop()
Ax0 += regParam * x0
U[I_start:I_end, :] = CG(
implicit_ATA(V, W, nomega, "U", use_MTTKRP), b, x0, r,
regParam, bsize, True)
else:
A = ctf.tensor((bsize, r, r))
t_form_EQs.start()
A.i("iuv") << V.i("Ju") * W.i("Ku") * nomega.i("iJK") * V.i(
"Jv") * W.i("Kv")
t_form_EQs.stop()
U[I_start:I_end, :] = CG(A, b, x0, r, regParam, bsize)
return U
if (string == "V"):
num_blocks = int((J + block_size - 1) / block_size)
for n in range(num_blocks):
J_start = n * block_size
J_end = min(J, J_start + block_size)
bsize = J_end - J_start
t_o_slice.start()
if num_blocks > 1:
nomega = omega[:, J_start:J_end, :]
else:
nomega = omega
t_o_slice.stop()
x0 = ctf.random.random((bsize, r))
b = ctf.tensor((bsize, r))
t_RHS.start()
if num_blocks == 1:
if use_MTTKRP:
ctf.MTTKRP(T, [U, b, W], 1)
else:
b.i("jr") << U.i("Ir") * W.i("Kr") * T.i("IjK") # RHS; ATb
else:
if use_MTTKRP:
ctf.MTTKRP(T[:, J_start:J_end, :], [U, b, W], 1)
else:
b.i("jr") << U.i("Ir") * W.i(
"Kr") * T[:, J_start:J_end, :].i("IjK") # RHS; ATb
t_RHS.stop()
if use_implicit:
Ax0 = ctf.tensor((bsize, r))
t_cg_TTTP.start()
if use_MTTKRP:
ctf.MTTKRP(ctf.TTTP(nomega, [U, x0, W]), [U, Ax0, W], 1)
else:
Ax0.i("jr") << U.i("Ir") * W.i("Kr") * ctf.TTTP(
nomega, [U, x0, W]).i("IjK")
t_cg_TTTP.stop()
Ax0 += regParam * x0
V[J_start:J_end, :] = CG(
implicit_ATA(U, W, nomega, "V", use_MTTKRP), b, x0, r,
regParam, bsize, True)
else:
A = ctf.tensor((bsize, r, r))
t_form_EQs.start()
A.i("juv") << U.i("Iu") * W.i("Ku") * nomega.i("IjK") * U.i(
"Iv") * W.i("Kv")
t_form_EQs.stop()
V[J_start:J_end, :] = CG(A, b, x0, r, regParam, bsize)
return V
if (string == "W"):
num_blocks = int((K + block_size - 1) / block_size)
for n in range(num_blocks):
K_start = n * block_size
K_end = min(K, K_start + block_size)
bsize = K_end - K_start
t_o_slice.start()
if num_blocks > 1:
nomega = omega[:, :, K_start:K_end]
else:
nomega = omega
t_o_slice.stop()
x0 = ctf.random.random((bsize, r))
b = ctf.tensor((bsize, r))
t_RHS.start()
if num_blocks == 1:
if use_MTTKRP:
ctf.MTTKRP(T, [U, V, b], 2)
else:
b.i("kr") << U.i("Ir") * V.i("Jr") * T.i("IJk") # RHS; ATb
else:
if use_MTTKRP:
ctf.MTTKRP(T[:, :, K_start:K_end], [U, V, b], 2)
else:
b.i("kr") << U.i("Ir") * V.i(
"Jr") * T[:, :, K_start:K_end].i("IJk") # RHS; ATb
t_RHS.stop()
if use_implicit:
Ax0 = ctf.tensor((bsize, r))
t_cg_TTTP.start()
if use_MTTKRP:
ctf.MTTKRP(ctf.TTTP(nomega, [U, V, x0]), [U, V, Ax0], 2)
else:
Ax0.i("kr") << U.i("Ir") * V.i("Jr") * ctf.TTTP(
nomega, [U, V, x0]).i("IJk")
t_cg_TTTP.stop()
Ax0 += regParam * x0
W[K_start:K_end, :] = CG(
implicit_ATA(U, V, nomega, "W", use_MTTKRP), b, x0, r,
regParam, bsize, True)
else:
A = ctf.tensor((bsize, r, r))
t_form_EQs.start()
A.i("kuv") << U.i(
"Iu") * V.i("Ju") * nomega.i("IJk") * U.i("Iv") * V.i(
"Jv") # LHS; ATA using matrix-vector multiplication
t_form_EQs.stop()
W[K_start:K_end, :] = CG(A, b, x0, r, regParam, bsize)
return W
def run_CCD(T,U,V,W,omega,regParam,num_iter,time_limit,objective_frequency,use_MTTKRP=True):
U_vec_list = []
V_vec_list = []
W_vec_list = []
r = U.shape[1]
for f in range(r):
U_vec_list.append(U[:,f])
V_vec_list.append(V[:,f])
W_vec_list.append(W[:,f])
# print(T)
# T.write_to_file('tensor_out.txt')
# assert(T.sp == 1)
ite = 0
objectives = []
t_before_loop = time.time()
t_obj_calc = 0.
t_CCD = ctf.timer_epoch("ccd_CCD")
t_CCD.begin()
while True:
t_iR_upd = ctf.timer("ccd_init_R_upd")
t_iR_upd.start()
t0 = time.time()
R = ctf.copy(T)
t1 = time.time()
# R -= ctf.einsum('ijk, ir, jr, kr -> ijk', omega, U, V, W)
R -= ctf.TTTP(omega, [U,V,W])
t2 = time.time()
# R += ctf.einsum('ijk, i, j, k -> ijk', omega, U[:,0], V[:,0], W[:,0])
R += ctf.TTTP(omega, [U[:,0], V[:,0], W[:,0]])
t3 = time.time()
t_iR_upd.stop()
t_b_obj = time.time()
if ite % objective_frequency == 0:
duration = time.time() - t_before_loop - t_obj_calc
[objective, RMSE] = get_objective(T,U,V,W,omega,regParam)
objectives.append(objective)
if glob_comm.rank() == 0:
print('Objective after',duration,'seconds (',ite,'iterations) is: {}'.format(objective))
print('RMSE after',duration,'seconds (',ite,'iterations) is: {}'.format(RMSE))
t_obj_calc += time.time() - t_b_obj
if glob_comm.rank() == 0 and status_prints == True:
print('ctf.copy() takes {}'.format(t1-t0))
print('ctf.TTTP() takes {}'.format(t2-t1))
print('ctf.TTTP() takes {}'.format(t3-t2))
for f in range(r):
# update U[:,f]
if glob_comm.rank() == 0 and status_prints == True:
print('updating U[:,{}]'.format(f))
t0 = time.time()
if use_MTTKRP:
alphas = ctf.tensor(R.shape[0])
#ctf.einsum('ijk -> i', ctf.TTTP(R, [None, V_vec_list[f], W_vec_list[f]]),out=alphas)
ctf.MTTKRP(R, [alphas, V_vec_list[f], W_vec_list[f]], 0)
else:
alphas = ctf.einsum('ijk, j, k -> i', R, V_vec_list[f], W_vec_list[f])
t1 = time.time()
if use_MTTKRP:
betas = ctf.tensor(R.shape[0])
#ctf.einsum('ijk -> i', ctf.TTTP(omega, [None, V_vec_list[f]*V_vec_list[f], W_vec_list[f]*W_vec_list[f]]),out=betas)
ctf.MTTKRP(omega, [betas, V_vec_list[f]*V_vec_list[f], W_vec_list[f]*W_vec_list[f]], 0)
else:
betas = ctf.einsum('ijk, j, j, k, k -> i', omega, V_vec_list[f], V_vec_list[f], W_vec_list[f], W_vec_list[f])
t2 = time.time()
U_vec_list[f] = alphas / (regParam + betas)
U[:,f] = U_vec_list[f]
if glob_comm.rank() == 0 and status_prints == True:
print('ctf.einsum() takes {}'.format(t1-t0))
print('ctf.einsum() takes {}'.format(t2-t1))
# update V[:,f]
if glob_comm.rank() == 0 and status_prints == True:
print('updating V[:,{}]'.format(f))
if use_MTTKRP:
alphas = ctf.tensor(R.shape[1])
#ctf.einsum('ijk -> j', ctf.TTTP(R, [U_vec_list[f], None, W_vec_list[f]]),out=alphas)
ctf.MTTKRP(R, [U_vec_list[f], alphas, W_vec_list[f]], 1)
else:
alphas = ctf.einsum('ijk, i, k -> j', R, U_vec_list[f], W_vec_list[f])
if use_MTTKRP:
betas = ctf.tensor(R.shape[1])
#ctf.einsum('ijk -> j', ctf.TTTP(omega, [U_vec_list[f]*U_vec_list[f], None, W_vec_list[f]*W_vec_list[f]]),out=betas)
ctf.MTTKRP(omega, [U_vec_list[f]*U_vec_list[f], betas, W_vec_list[f]*W_vec_list[f]], 1)
else:
betas = ctf.einsum('ijk, i, i, k, k -> j', omega, U_vec_list[f], U_vec_list[f], W_vec_list[f], W_vec_list[f])
V_vec_list[f] = alphas / (regParam + betas)
V[:,f] = V_vec_list[f]
if glob_comm.rank() == 0 and status_prints == True:
print('updating W[:,{}]'.format(f))
if use_MTTKRP:
alphas = ctf.tensor(R.shape[2])
#ctf.einsum('ijk -> k', ctf.TTTP(R, [U_vec_list[f], V_vec_list[f], None]),out=alphas)
ctf.MTTKRP(R, [U_vec_list[f], V_vec_list[f], alphas], 2)
else:
alphas = ctf.einsum('ijk, i, j -> k', R, U_vec_list[f], V_vec_list[f])
if use_MTTKRP:
betas = ctf.tensor(R.shape[2])
#ctf.einsum('ijk -> k', ctf.TTTP(omega, [U_vec_list[f]*U_vec_list[f], V_vec_list[f]*V_vec_list[f], None]),out=betas)
ctf.MTTKRP(omega, [U_vec_list[f]*U_vec_list[f], V_vec_list[f]*V_vec_list[f], betas], 2)
else:
betas = ctf.einsum('ijk, i, i, j, j -> k', omega, U_vec_list[f], U_vec_list[f], V_vec_list[f], V_vec_list[f])
W_vec_list[f] = alphas / (regParam + betas)
W[:,f] = W_vec_list[f]
t_tttp = ctf.timer("ccd_TTTP")
t_tttp.start()
R -= ctf.TTTP(omega, [U_vec_list[f], V_vec_list[f], W_vec_list[f]])
if f+1 < r:
R += ctf.TTTP(omega, [U_vec_list[f+1], V_vec_list[f+1], W_vec_list[f+1]])
t_tttp.stop()
t_iR_upd.stop()
ite += 1
if ite == num_iter or time.time() - t_before_loop - t_obj_calc > time_limit:
break
t_CCD.end()
duration = time.time() - t_before_loop - t_obj_calc
[objective, RMSE] = get_objective(T,U,V,W,omega,regParam)
if glob_comm.rank() == 0:
print('CCD amortized seconds per sweep: {}'.format(duration/ite))
print('Time/CCD Iteration: {}'.format(duration/ite))
print('Objective after',duration,'seconds (',ite,'iterations) is: {}'.format(objective))
print('RMSE after',duration,'seconds (',ite,'iterations) is: {}'.format(RMSE))