Processing time problem

[dasda-asd]

I plan to use this code to process the images in the dataset to get the appropriate depth map. The dataset has 365 images. But it takes about two seconds to process each image. I guess this is because I have to load the model every time I process an image.
Could you please help me out there?

Code:
import torch
from PIL import Image
import numpy as np
from vadepthnet.networks.vadepthnet import VADepthNet
from vadepthnet.dataloaders.dataloader import ToTensor
import os
import tqdm
import glob
import time

def save_raw_16bit(depth, fpath="raw.png"):
if isinstance(depth, torch.Tensor):
depth = depth.squeeze().cpu().numpy()
assert isinstance(depth, np.ndarray), "Depth must be a torch tensor or numpy array"
assert depth.ndim == 2, "Depth must be 2D"
depth = depth * 5000 # scale for 16-bit png
depth = depth.astype(np.uint16)
depth = Image.fromarray(depth)
depth.save(fpath)
print("Saved raw depth to", fpath)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device: %s" % device)
model = VADepthNet(max_depth=10,
prior_mean=1.54,
img_size=(480, 640))
model = torch.nn.DataParallel(model)
checkpoint = torch.load('vadepthnet_nyu.pth', map_location=device)
model.load_state_dict(checkpoint['model'])
model.eval()
totensor = ToTensor('test')

#Dataset path
#There are 365 images in the dataset.
img_dir = 'D:/xiaohe/VA-depthnet/VA-DepthNet-main/rgb/*.png'
for img in glob.glob(img_dir):
start = time.perf_counter()
img = 'D:/xiaohe/VA-depthnet/VA-DepthNet-main/rgb/' + os.path.basename(img)
image = Image.open(img)
image = np.asarray(image, dtype=np.float32) / 255.0
image = totensor.to_tensor(image)
image = totensor.normalize(image)
image = image.unsqueeze(0)
pdepth = model(image)
pdepth = pdepth.cpu().detach().numpy()
pdepth = np.squeeze(pdepth)
fpath = "./output/"
fpath = os.path.join(fpath, os.path.basename(img))
save_raw_16bit(pdepth, fpath)
end = time.perf_counter()
runTime = end - start
runTime_ms = runTime * 1000
print(runTime, "s")
print(runTime_ms, "ms")

Result:
D:\xiaohe\anaconda\envs\zoedepth\python.exe D:\xiaohe\VA-depthnet\VA-DepthNet-main\test-author.py
device: cuda
D:\xiaohe\anaconda\envs\zoedepth\lib\site-packages\torch\functional.py:504: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at C:\cb\pytorch_1000000000000\work\aten\src\ATen\native\TensorShape.cpp:3191.)
return _VF.meshgrid(tensors, **kwargs) # type: ignore[attr-defined]
Saved raw depth to ./output/1341845820.751833.png
3.3062937000000003 s
3306.2937 ms
Saved raw depth to ./output/1341845820.787768.png
2.1543924000000008 s
2154.3924000000006 ms
Saved raw depth to ./output/1341845820.819654.png
2.1151333 s
2115.1333 ms
Saved raw depth to ./output/1341845820.851997.png
2.1185686000000015 s
2118.5686000000014 ms
Saved raw depth to ./output/1341845820.887882.png
2.2202800999999983 s
2220.280099999998 ms
Saved raw depth to ./output/1341845820.920082.png
2.1999916000000006 s
2199.9916000000007 ms
Traceback (most recent call last):
File "D:\xiaohe\VA-depthnet\VA-DepthNet-main\test-author.py", line 49, in
pdepth = model(image)
File "D:\xiaohe\anaconda\envs\zoedepth\lib\site-packages\torch\nn\modules\module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "D:\xiaohe\anaconda\envs\zoedepth\lib\site-packages\torch\nn\parallel\data_parallel.py", line 169, in forward
return self.module(*inputs[0], **kwargs[0])
File "D:\xiaohe\anaconda\envs\zoedepth\lib\site-packages\torch\nn\modules\module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "D:\xiaohe\VA-depthnet\VA-DepthNet-main\vadepthnet\networks\vadepthnet.py", line 303, in forward
d = self.vlayer(x)
File "D:\xiaohe\anaconda\envs\zoedepth\lib\site-packages\torch\nn\modules\module.py", line 1190, in _call_impl
return forward_call(*input, **kwargs)
File "D:\xiaohe\VA-depthnet\VA-DepthNet-main\vadepthnet\networks\vadepthnet.py", line 188, in forward
x = torch.linalg.solve(ATA+jitter,ATB)
KeyboardInterrupt.

[cnexah]

1. Your time includes the time for loading the image, and saving the predicted depth map.
2. I suggest to upload the image to GPU for processing. For now the model is making prediction on CPU.

[cnexah]

The above code use 'cpu' as device. Could you try the following code using gpu?

import torch
from PIL import Image
import numpy as np
from vadepthnet.networks.vadepthnet import VADepthNet
from vadepthnet.dataloaders.dataloader import ToTensor

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("device: %s" % device)

model = VADepthNet(max_depth=10,
prior_mean=1.54,
img_size=(480, 640))
model = torch.nn.DataParallel(model).cuda()
checkpoint = torch.load('vadepthnet_nyu.pth', map_location=device)
model.load_state_dict(checkpoint['model'])
model.eval()
img = Image.open(image_path)
img = np.asarray(img, dtype=np.float32) / 255.0
#img = torch.from_numpy(img).cuda().unsqueeze(0)

totensor = ToTensor('test')
img = totensor.to_tensor(img)
img = totensor.normalize(img)
img = img.unsqueeze(0).cuda()
pdepth = model.forward(img)
print(pdepth)