Suggestions for improving training speed (especially when input data is large)

[Robotwithsoul]

Hi, In memory.py, I suggested to change a little bit of your codes, which will be helpful for improving the training speed (especially when the input image data is 3D)

• your original codes

state = torch.stack([trans.state for trans in transition[:self.history]]).to(dtype=torch.float32, device=self.device).div_(255)
next_state = torch.stack([trans.state for trans in transition[self.n:self.n + self.history]]).to(dtype=torch.float32, device=self.device).div_(255)

• suggested codes

state = torch.stack([trans.state for trans in transition[:self.history]]).to(device=self.device).to(dtype=torch.float32).div_(255)
next_state = torch.stack([trans.state for trans in transition[self.n:self.n + self.history]]).to(device=self.device).to(dtype=torch.float32).div_(255)

Here is the code for testing：

import timeit
import numpy as np
import torch

T,T1=[],[]

device=torch.device('cuda')
for i in range (0,4):

    A=np.zeros((100,100,100),dtype=np.int)
    B=torch.tensor(A,dtype=torch.int)
    T.append(B)
for i in range (0,4):
    A=np.zeros((100,100,100),dtype=np.int)
    B=torch.tensor(A,dtype=torch.int)
    T1.append(B)

# This line is used for initilization
M=torch.stack(T).to(dtype=torch.float32, device=device).div_(255)

# Comparison 
timea=timeit.default_timer()
M=torch.stack(T).to(dtype=torch.float32, device=device).div_(255)
timeb=timeit.default_timer()
N=torch.stack(T1).to(device=device).to(dtype=torch.float32).div_(255)
timec=timeit.default_timer()

print("time1 is:{}\n time2 is:{}".format(timeb-timea,timec-timeb))

[Kaixhin]

How odd - do you have any idea why this is the case? Thanks a lot for the script for quick testing - can confirm that I am indeed seeing a speedup (not so significant for smaller tensors, but sure). Feel free to submit a PR for this change!

[Robotwithsoul]

I'm not sure, but I guess this is because that the following codes perform the data type conversion in CPU

.to(dtype=torch.float32, device=device)

while the suggested codes perform the data type conversion in GPU

.to(device=device).to(dtype=torch.float32)