Hi,
I’m new to pytorch, and I’m doing something a little heterodox. I build a very dynamic computational graph with a lot of inplace operations due to the complex nature of my subject. I expect it to be slow as I build the whole ‘network’ from scratch, but the backward is just so slow (about 1 minutes for forward and 20+ minutes for backward) that I must have done something wrong. So I’m asking for some advice on implementing pytorch for general purpose computation. Below is an example of my code:
First, I defined a function to update the state of one node:
def neuronForward(self, InputS, InputR, InputTau, StateTau, endTime, V0, Dv0, G, Tau_m, H, Tau_n, I, Tau_r, Tau_s, G_exc, G_inh ):
	Delta = (Tau_m+G*Tau_n)**2 - 4*Tau_m*Tau_n*(H+G)
	deta = Delta.data[0]
	nptr = (-1/Tau_r).data[0]
	npts = (-1/Tau_s).data[0]
	P = G/Tau_m +1/Tau_n
	Q = (H+G)/(Tau_m*Tau_n)
	L = I/(H+G)
	Cr = 1/(Tau_m*Tau_n) - 1/(Tau_m*Tau_r)
	Cs = 1/(Tau_m*Tau_n) - 1/(Tau_m*Tau_s)
	PQR = 1-P*Tau_r+Q*Tau_r**2
	PQS = 1-P*Tau_s+Q*Tau_s**2

	InheritMask = (InputTau == StateTau).float().detach()
	EffectiveMask = (InputTau > 0).float().detach()
	GexcMask = (InputR > 0).float().detach()
	GinhMask = (InputR < 0).float().detach()

	S0 = torch.sum(InheritMask*InputS)
	Kr = EffectiveMask*(GexcMask*G_exc*Tau_r*InputR/(Tau_r-Tau_s) - GinhMask*G_inh*Tau_r*InputR/(Tau_r-Tau_s))
	Ks = EffectiveMask*(InputS - GexcMask*G_exc*Tau_r*InputR/(Tau_r-Tau_s) + GinhMask*G_inh*Tau_r*InputR/(Tau_r-Tau_s))

	if(deta >0):
		r1 = -0.5/Tau_n -0.5*G/Tau_m +Delta**0.5/(2*Tau_m*Tau_n)
		r2 = -0.5/Tau_n -0.5*G/Tau_m -Delta**0.5/(2*Tau_m*Tau_n)
		if(r1.data[0]!= nptr and r2.data[0]!= nptr):
			Vr = Cr*Kr*Tau_r**2*torch.exp(-InputTau/Tau_r)/PQR
			C1r = Kr*(1/(Tau_m*(r1-r2)) + Cr*(Tau_r+r2*Tau_r**2)/((r1-r2)*PQR))
			C2r = -Cr*Kr*Tau_r**2/PQR-C1r
			Vr = Vr + C1r*torch.exp(r1*InputTau) + C2r*torch.exp(r2*InputTau)
			Dvr = -Cr*Kr*Tau_r*torch.exp(-InputTau/Tau_r)/PQR + C1r*r1*torch.exp(r1*InputTau) + C2r*r2*torch.exp(r2*InputTau)
		else:
			Vr = Cr*Kr*Tau_r*InputTau*torch.exp(-InputTau/Tau_r)/(P*Tau_r-2)
			C1r = Kr*(Tau_r*(P-Cr*Tau_m)-2)/((r1-r2)*(P*Tau_r-2)*Tau_m)
			C2r = -C1r
			Vr = Vr + C1r*torch.exp(r1*InputTau) + C2r*torch.exp(r2*InputTau)
			Dvr = Cr*Kr*Tau_r*torch.exp(-InputTau/Tau_r)/(P*Tau_r - 2) - Cr*Kr*InputTau*torch.exp(-InputTau/Tau_r)/(P*Tau_r - 2) \
				+ C1r*r1*torch.exp(r1*InputTau) + C2r*r2*torch.exp(r2*InputTau)
		if(r1.data[0]!= npts and r2.data[0]!= npts):
			Vs = Cs*Ks*Tau_s ** 2*torch.exp(-InputTau/Tau_s)/PQS
			C1s = Ks*(1/(Tau_m*(r1 - r2)) + Cs*(Tau_s + r2*Tau_s ** 2)/((r1 - r2)*PQS))
			C2s = -Cs*Ks*Tau_s ** 2/PQS - C1s
			Vs = Vs + C1s*torch.exp(r1*InputTau) + C2s*torch.exp(r2*InputTau)
			Dvs = -Cs*Ks*Tau_s*torch.exp(-InputTau/Tau_s)/PQS + C1s*r1*torch.exp(r1*InputTau) + C2s*r2*torch.exp(r2*InputTau)
		else:
			Vs = Cs*Ks*Tau_s*InputTau*torch.exp(-InputTau/Tau_s)/(P*Tau_s-2)
			C1s = Ks*(Tau_s*(P-Cs*Tau_m)-2)/((r1-r2)*(P*Tau_s-2)*Tau_m)
			C2s = -C1s
			Vs = Vs + C1s*torch.exp(r1*InputTau) + C2s*torch.exp(r2*InputTau)
			Dvs = Cs*Ks*Tau_s*torch.exp(-InputTau/Tau_s)/(P*Tau_s - 2) - Cs*Ks*InputTau*torch.exp(-InputTau/Tau_s)/(P*Tau_s - 2) \
					+ C1s*r1*torch.exp(r1*InputTau) + C2s*r2*torch.exp(r2*InputTau)
		C3 = (Dv0-S0/Tau_m+r2*L-r2*V0)/(r1-r2)
		C4 = V0-C3-L
		VI = C3*torch.exp(r1*StateTau)+C4*torch.exp(r2*StateTau)+L
		DvI = C3*r1*torch.exp(r1*StateTau) + C4*r2*torch.exp(r2*StateTau)
	elif(deta < 0):
		alpha = -0.5/Tau_n -0.5*G/Tau_m
		beta = (-Delta)**0.5/(2*Tau_m*Tau_n)
		Vr = Cr*Kr*Tau_r ** 2*torch.exp(-InputTau/Tau_r)/PQR
		C1r = -Cr*Kr*Tau_r**2/PQR
		C2r = Cr*Kr*Tau_r*(1+alpha*Tau_r)/(beta*PQR)+Kr/(beta*Tau_m)
		Vr = Vr + torch.exp(alpha*InputTau)*(C1r*torch.cos(beta*InputTau)+C2r*torch.sin(beta*InputTau))
		Vs = Cs*Ks*Tau_s ** 2*torch.exp(-InputTau/Tau_s)/PQS
		C1s = -Cs*Ks*Tau_s**2/PQS
		C2s = Cs*Ks*Tau_s*(1+alpha*Tau_s)/(beta*PQS)+Ks/(beta*Tau_m)
		Vs = Vs + torch.exp(alpha*InputTau)*(C1s*torch.cos(beta*InputTau)+C2s*torch.sin(beta*InputTau))
		C3 = V0 - L
		C4 = (Dv0 - S0/Tau_m -alpha*C3)/beta
		VI = torch.exp(alpha*StateTau)*(C3*torch.cos(beta*StateTau)+C4*torch.sin(beta*StateTau)) + L
		Dvr = -Cr*Kr*Tau_r*torch.exp(-InputTau/Tau_r)/PQR + torch.exp(alpha*InputTau)*((alpha*C1r+beta*C2r)*torch.cos(beta*InputTau) +
		                                       (alpha*C2r-beta*C1r)*torch.sin(beta*InputTau))
		Dvs = -Cs*Ks*Tau_s*torch.exp(-InputTau/Tau_s)/PQS + torch.exp(alpha*InputTau)*((alpha*C1s+beta*C2s)*torch.cos(beta*InputTau) +
		                                       (alpha*C2s-beta*C1s)*torch.sin(beta*InputTau))
		DvI = torch.exp(alpha*StateTau)*((alpha*C3 + beta*C4)*torch.cos(beta*StateTau) +
		                                 (alpha*C4 - beta*C3)*torch.sin(beta*StateTau))
	elif(deta==0):
		r = -0.5/Tau_n -0.5*G/Tau_m
		if(r!=nptr):
			Vr = Cr*Kr*Tau_r**2*torch.exp(-InputTau/Tau_r)/PQR
			C1r = -Cr*Kr*Tau_r**2/PQR
			C2r = Cr*Kr*Tau_r*(1+r*Tau_r)/PQR +Kr/Tau_m
			Vr = Vr + (C1r+C2r*InputTau)*torch.exp(r*InputTau)
			Dvr = -Cr*Kr*Tau_r*torch.exp(-InputTau/Tau_r)/PQR + (r*C1r + (r*InputTau + 1)*C2r)*torch.exp(r*InputTau)
		else:
			Vr = Cr*Kr*InputTau**2*torch.exp(-InputTau/Tau_r)/2
			C2r = Kr/Tau_m
			Vr = Vr + C2r*InputTau*torch.exp(r*InputTau)
			Dvr = Cr*Kr*InputTau*torch.exp(-InputTau/Tau_r) - Cr*Kr*InputTau ** 2*torch.exp(-InputTau/Tau_r)/(2*Tau_r)+\
			      (r*InputTau+1)*C2r*torch.exp(r*InputTau)
		if(r!=npts):
			Vs = Cs*Ks*Tau_s ** 2*torch.exp(-InputTau/Tau_s)/PQS
			C1s = -Cs*Ks*Tau_s**2/PQS
			C2s = Cs*Ks*Tau_s*(1+r*Tau_s)/PQS +Ks/Tau_m
			Vs = Vs + (C1s+C2s*InputTau)*torch.exp(r*InputTau)
			Dvs = -Cs*Ks*Tau_s*torch.exp(-InputTau/Tau_s)/PQS + (r*C1s + (r*InputTau + 1)*C2s)*torch.exp(r*InputTau)
		else:
			Vs = Cs*Ks*InputTau**2*torch.exp(-InputTau/Tau_s)
			C2s = Ks/Tau_m
			Vs = Vs + C2s*InputTau*torch.exp(r*InputTau)
			Dvs = Cs*Ks*InputTau*torch.exp(-InputTau/Tau_s) - Cs*Ks*InputTau ** 2*torch.exp(-
	V = torch.sum(Vr+Vs) + VI
	Dv = torch.sum(Dvr + Dvs) + DvI
	endTime = couple(V, endTime, -Dv.data[0])
	theTime = 0*InputTau + endTime
	InputTau = couple(theTime,InputTau,1)
	StateTau = couple(endTime,StateTau,1)

	InputS = Kr*(torch.exp(-InputTau/Tau_r)-torch.exp(-InputTau/Tau_s)) + EffectiveMask*InputS*torch.exp(-InputTau/Tau_s) +\
	         (1-EffectiveMask)*InputS
	InputR = EffectiveMask*InputR*torch.exp(-InputTau/Tau_r) + (1-EffectiveMask)*InputR

	return V, endTime, Dv, InputR, InputS

And then, I update the state one-by-one in the main loop:
for ii in range(StateUpdateNum):
	neuronInd = int(SrIndex[ii])
	StateTau = SrTime[ii] - Time[neuronInd]
	InpTau = SrTime[ii] - PreSynpticTimeList[neuronInd]
	InpTau = torch.min(InpTau, StateTau)
	endTime = Variable(torch.Tensor([SrTime[ii]]), requires_grad=True)
	theV, endTime, theDv, InputR, InputS = \
		neuronForward(S[neuronInd], R[neuronInd], InpTau, StateTau, endTime, V[neuronInd], Dv[neuronInd], theG,
					  Tau_m[neuronInd],
					  H[neuronInd], theTau_n, theI, Tau_r[neuronInd], Tau_s[neuronInd], theG_exc, theG_inh)

	Time[neuronInd] = endTime
	S[neuronInd] = InputS
	R[neuronInd] = InputR
	Dv = Dv.clone()
	Dv[neuronInd] = theDv
	V = V.clone()
	V[neuronInd] = theV.detach()
	State[neuronInd] = SrState[ii]
	if (SrState[ii] == 4):
		if (NetSpkTime is None):
			NetSpkTime = endTime.clone()
		else:
			NetSpkTime = torch.cat((NetSpkTime, endTime))
		for postInd in ForwardLink[neuronInd + InpNeuronNum]:
			if (len(PreSynpticTimeList[postInd]) == 0):
				PreSynpticTimeList[postInd] = preTime + DelaysMatrix[postInd, neuronInd + InpNeuronNum]
			else:
				PreSynpticTimeList[postInd] = torch.cat(
					(PreSynpticTimeList[postInd], preTime + DelaysMatrix[postInd, neuronInd + InpNeuronNum]))
			if len(S[postInd]) == 0:
				S[postInd] = Variable(torch.zeros(1), requires_grad=True).clone()
				R[postInd] = WeightsMatrix[postInd, neuronInd + InpNeuronNum].clone()
			else:
				S[postInd] = torch.cat((S[postInd], Variable(torch.zeros(1), requires_grad=True)))
				R[postInd] = torch.cat((R[postInd], WeightsMatrix[postInd, neuronInd + InpNeuronNum]))

I understand that my ‘network’ is too fragmentary, and is really bad for vectorization. But the ‘network’ is so dynamic that I can’t know each state and connectivity before state update. Even through, it shouldn’t be that slow. I can do something to speed it up? Thanks in advance. And by the way, I can’t find the module torch.autograd.profiler in my pytorch (linux python2.7 conda installation latest version), do I need a seperate installation?