My program’s memory usage is roughly an order of magnitude greater when I specify requires_grad=True on the parameters of my model. I’ve looked through the docs to find a way to reduce my program’s memory consumption, but I can’t seem to figure it out. Here is my objective function:
def fun(x, cons, est, trans, model, data):
print(x)
for con in cons:
valid = np.all(con['fun'](x, *con['args']) > 0)
if valid == False: break
if valid == False:
ll = float('nan')
grad = np.empty(len(x))
grad[:] = np.nan
print('Constraint violation')
print()
return (-ll, -grad)
else:
torch.manual_seed(7)
ex0 = (None,)
ex1 = (slice(None), None)
x = torch.tensor(x.astype(np.float32), requires_grad=True)
ll = 0
for d in range(len(data)):
N, S, K = (data[d]['N'], data[d]['S'], data[d]['K'])
p = 0
mu = torch.tensor([])
re = []
theta = {}
for (key, val) in est.items():
if type(val) != str:
theta[key] = torch.tensor(val)
elif val == 'pool':
theta[key] = x[p][ex0]
p += 1
elif val == 'rand':
mu = torch.cat((mu, x[p][ex0]))
p += 1
re.append(key)
R = len(re)
di = tuple([torch.tensor(ind, dtype=torch.long)
for ind in np.diag_indices(R)])
tril = tuple([torch.tensor(ind, dtype=torch.long)
for ind in np.tril_indices(R, -1)])
triu = tuple([torch.tensor(ind, dtype=torch.long)
for ind in np.triu_indices(R, 1)])
sigma = torch.empty(R, R)
sigma[di] = x[p:p+R]
sigma[tril] = x[p+R:]
sigma[triu] = sigma.t()[triu]
sample = MultivariateNormal(mu, sigma).rsample((N * S,))
theta.update({re[r]: trans[re[r]](sample[:,r]) for r in range(R)})
theta['upsilon'] = torch.zeros(N * S)[ex1]
if K > 1:
theta['upsilon'] = torch.cat((theta['upsilon'],
theta['upsilon2'][ex1]), 1)
if K > 2:
theta['upsilon'] = torch.cat((theta['upsilon'],
theta['upsilon3'][ex1]), 1)
theta['upsilon'] = theta['upsilon'] / theta['kappa']
if data[d]['treat'] == 'pvd': theta['omega'] = torch.zeros(1)
ll = ll + logL(theta, model, data[d])
ll = ll / 1e+2
ll.backward()
grad = np.array(x.grad.cpu())
print(ll.item() * 1e+2)
print()
return (-ll.item(), -grad)
Here is the logL function referenced in the objective function:
def logL(theta, model, data):
N, S, T = (data['N'], data['S'], data['T'])
p_a = torch.ones(N * S)
model.start(theta, data['mat'])
for t in range(1, T + 1):
model.update(data, t)
p_a = p_a * (npeat(data['obs'][:,:,t-1], (S, 1)) * model.p_a).sum(1)
ll = p_a.reshape(N, S).mean(1).log().sum()
return ll
And here is my model
class Model:
def __init__(self, opt):
fun = {'p_k': {'rand': Probability.rand,
'hmax': Probability.hmax,
'smax': Probability.smax},
'p_ak': {'rand': Probability.rand,
'hmax': Probability.hmax,
'smax': Probability.smax},
'U' : {'none': Utility.Learning.none,
'diff': Utility.Learning.diff},
'v': {'lin': Utility.lin},
'w': {'lin': Probability.Weighting.lin,
'linlog': Probability.Weighting.linlog,
'ratio': Probability.Weighting.ratio},
'b_1': {'rand': Probability.rand,
'hmax': Probability.hmax,
'smax': Probability.smax},
'b_k': {'rand': Probability.rand,
'hmax': Probability.hmax,
'smax': Probability.smax},
'B': {'none': Activation.Base.none,
'exp': Activation.Base.exp,
'pow': Activation.Base.pow},
'S': {'none': Activation.Associative.none,
'chain': Activation.Associative.chain},
'P': {'none': Activation.Matching.none,
'part': Activation.Matching.part},
'R_a': {'lin': Utility.lin},
'R_c': {'none': Activation.Base.Learning.none,
'belief': Activation.Base.Learning.belief,
'reinf': Activation.Base.Learning.reinf,
'attrac': Activation.Base.Learning.attrac},
'S_q': {'none': Activation.Associative.Learning.none,
'diff': Activation.Associative.Learning.diff}}
self.fun = {key: fun[key][val] for (key, val) in opt.items()}
self.opt = opt
def start(self, theta, mat):
ex1 = (slice(None), None)
ex12 = (slice(None), None, None)
omega = theta['omega'][ex1].expand(-1, mat.size()[-1])
W = torch.cat((omega, 1 - omega), 1)
xi = W * theta['xi'][ex1]
beta = W * theta['beta'][ex1]
arg = {'p_k': {'rand': (1,),
'hmax': (1,),
'smax': (theta['kappa'][ex1], 1)},
'p_ak': {'rand': (1,),
'hmax': (1,),
'smax': (theta['lam'][ex12], 1)},
'U': {'none': (),
'diff': (theta['alpha'][ex1],)},
'v': {'lin': (mat,)},
'w': {'lin': (),
'linlog': (theta['gamma'][ex12], theta['delta'][ex12]),
'ratio': (theta['gamma'][ex12], theta['delta'][ex12], 2)},
'b_1': {'rand': (2,),
'hmax': (2,),
'smax': (1, 2)},
'b_k': {'rand': (2,),
'hmax': (2,),
'smax': (theta['iota'][ex12], 2)},
'B': {'cons': {},
'exp': {'gamma': theta['phi'][ex1]},
'pow': {'d': theta['phi'][ex12], 'dim': 2}},
'S': {'none': (),
'chain': (2,)},
'P': {'none': (),
'part': (xi[ex1], 2)},
'R_a': {'lin': ()},
'R_c': {'none': {},
'belief': {},
'reinf': {},
'attrac': {'delta': theta['rho'][ex1]}},
'S_q': {'none': (),
'diff': (theta['psi'][ex12], beta[ex1])}}
self.arg = {key: arg[key][val] for (key, val) in self.opt.items()}
self.theta = theta
self.p_ak = torch.zeros(1)
self.U = theta['upsilon']
self.v = self.fun['v'](*self.arg['v'])
self.B = torch.zeros(1).log()
self.S = torch.zeros(1)[ex1]
self.R_ct = torch.zeros(0)
self.S_q = torch.zeros(1)
def update(self, dat, t):
old = (Ellipsis, t - 1)
new = (Ellipsis, t)
ex1 = (slice(None), None)
ex2 = (slice(None), slice(None), None)
N, S, A, X, M, C, K, I = (dat['N'], dat['S'], dat['A'], dat['X'],
dat['M'], dat['C'], dat['K'], dat['I'])
hypo, hist, act, cue, sim, pay = (dat['hypo'], dat['hist'], dat['act'],
dat['cue'], dat['sim'], dat['pay'])
hypo_old = npeat(hypo[old], (S,)+(1,)*(hypo[old].dim()-1))
hist_old = npeat(hist[old], (S,)+(1,)*(hist[old].dim()-1))
act_old = npeat(act[old], (S,)+(1,)*(act[old].dim()-1))
cue_old = npeat(cue[old], (S,)+(1,)*(cue[old].dim()-1))
cue_new = npeat(cue[new], (S,)+(1,)*(cue[new].dim()-1))
sim_new = npeat(sim[new], (S,)+(1,)*(sim[new].dim()-1))
pay_old = npeat(pay[old], (S,)+(1,)*(pay[old].dim()-1))
fun, arg = (self.fun, self.arg)
R_a = fun['R_a'](pay_old, *arg['R_a'])
R_k = (self.p_ak * R_a[ex2]).sum(1)
self.U = fun['U'](self.U, R_k, int(t > 1), *arg['U'])
attn = 1 if t > 1 else self.theta['pi']
R_c = fun['R_c'](attn, hypo=hypo_old, hist=hist_old, **arg['R_c'])
self.R_ct = torch.cat((self.R_ct, R_c[ex2]), 2)
R_ct = self.R_ct[:,:,:t]
time = torch.arange(float(t), 0, -1)
self.S_q = fun['S_q'](self.S_q, self.S[ex2],
act_old[ex2], cue_old, *arg['S_q'])
self.B = fun['B'](B=self.B, t=time, R=R_c, R_j=R_ct, **arg['B'])
self.S = fun['S'](self.S_q, cue_new, *arg['S'])
P = fun['P'](sim_new, *arg['P'])
A_c = (self.B + self.S + P).reshape(N * S, I * M, C // M)
b_1 = fun['b_1'](A_c, *arg['b_1']).reshape(N * S, I * M, A, X).sum(3)
w = npeat(fun['w'](b_1, *arg['w']), (1, A // M, 1))
V_k = (w * self.v).sum(2)[ex2]
for k in range(1, K):
own = (slice(None), slice(0, A), k - 1)
opp = (slice(None), slice(A, 2 * A), k - 1)
V = torch.cat((V_k[opp], V_k[own]), 1).reshape(N * S, 2, A)
b_k = npeat(fun['b_k'](V, *arg['b_k']), (1, A, 1))
V_k = torch.cat((V_k, (b_k * self.v).sum(2)[ex2]), 2)
self.p_k = fun['p_k'](self.U, *arg['p_k'])
self.p_ak = fun['p_ak'](V_k[:,:A], *arg['p_ak'])
self.p_a = (self.p_k[ex1] * self.p_ak).sum(2)
I don’t understand why autograd is consuming so much memory. Please let me know if you see anything I can do to reduce my memory consumption. I’m willing to provide any additional information related to the program if you think it will help diagnose the problem.
