I am working with a pytorch based code. Part of my code has 4 nested for loops. Basically, it is finding patches of image and estimating similarity between two patches. Since the parameters are torch elements, none of the python libraries like joblib are working. I am a newbie to pycuda and I kindly seek help to parallelize this code. At present, this is slow computation that takes more than 1.5 seconds. Here is part of my code.
import torch
import numpy as np
import cv2
import time
def mul_val(a,b,ax=None):
return torch.mean(((b-a)/(b+0.01))**2)
def np_to_torch(img_np):
return torch.from_numpy(img_np)[None,None, :]
def main_fn(a,alpha = 0.25):
start = time.time()
final_u = torch.zeros(w,h)
frame1 = a
for y1 in range(w):
i = y1*grid_size
for x1 in range(h):
j = x1*grid_size ## batch, channel, width, height
block1 = frame1[:,:,i:i+grid_size, j:j+grid_size]
corr_list = []
for y2 in range(y1-radius,y1+radius+1):
i2 = y2*grid_size
if not (0 <= y2 < h):
continue
for x2 in range(x1-radius,x1+radius+1):
j2 = x2*grid_size
if not (0 <= x2 < w):
continue
block2 = frame1[:,:,i2:i2+grid_size, j2:j2+grid_size]
if not (block1.shape == block2.shape):
continue
corr = mul_val(block1, block2)
corr_list.append(corr)
corr_list=sorted(corr_list, reverse = True)
del corr_list[20:]
uncorr = 1.0 - (sum(corr_list)/20.0)
final_u[y1,x1] = torch.mul(a[:,:,i,j],uncorr)
new_val = torch.norm(torch.sub(1.0,final_u))
final_value = alpha*new_val
print("Time taken is ..", time.time()-start)
return final_value
grid_size = 9
radius = 3
im1 = cv2.imread("1.png", 0)
img = np.asarray(im1)
img = img.astype(np.float64)
img_torch = np_to_torch(img) ### made this as a torch element intentionally.
frame_width = img.shape[0]
frame_height = img.shape[1]
h = int(frame_height//grid_size)
w = int(frame_width//grid_size)
main_fn(img_torch)
The above code opens an image and converts to torch element just to replicate the parameters as torch elements. Any help in this regard is highly appreciated. Thank you