Conv

Scenario Description

Conv2d: used for edge detection and feature extraction.

Conv3d: used for spatiotemporal feature extraction.

Currently, KuDNN supports the torch.float16 and torch.float32 data types. For other data types, see the open-source branch.

Sample Code

import torch
import torch.nn as nn

# Enable KuDNN.
torch._C._set_kdnn_enabled(True)

# Conv2d example
conv2d = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1)
x = torch.randn(8, 3, 32, 32) # [batch, channel, H, W] # Defaul to fp32
y = conv2d (x) # Output: [8, 16, 30, 30]
print ("Conv2d output shape", y.shape)
print(y)

# Conv3d example
conv3d = nn.Conv3d(1, 8, kernel_size=(3,3,3)) 
x = torch.randn(4, 1, 10, 64, 64) # [batch, channel, depth, H, W] # Defaul to fp32
y = conv3d (x) # Output: [4, 8, 8, 62, 62]

print ("Conv3d output shape", y.shape)
print(y)

Parent topic: Examples