[Python] Doing PyWavelet calculation on GPU

Currently working on a classifier using PyWavelet, here is my calculation block:

class WaveletLayer(nn.Module):
def __init__(self):
super(WaveletLayer, self).__init__()

def forward(self, x):
def wavelet_transform(img):
coeffs = pywt.dwt2(img.cpu().numpy(), "haar")
LL, (LH, HL, HH) = coeffs
return (
torch.from_numpy(LL).to(img.device),
torch.from_numpy(LH).to(img.device),
torch.from_numpy(HL).to(img.device),
torch.from_numpy(HH).to(img.device),
)

# Apply wavelet transform to each channel separately
LL, LH, HL, HH = zip(
*[wavelet_transform(x[:, i : i + 1]) for i in range(x.shape[1])]
)

# Concatenate the results
LL = torch.cat(LL, dim=1)
LH = torch.cat(LH, dim=1)
HL = torch.cat(HL, dim=1)
HH = torch.cat(HH, dim=1)

return torch.cat([LL, LH, HL, HH], dim=1)



The output from this module goes to a resnet block for learning, while doing this I find my CPU clogged and thus slowing down my training process

I am trying to use the GPUs for these calculations.

Continue reading...