Linear

Scenario Description

Perform linear transformation on the input tensor. Currently, KuDNN supports the torch.int8, 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)

# Linear example
# Input data: (batch, in_features). The default data type is torch.float32.
input_tensor = torch.randn(128, 20)

# Construct the Linear layer.
linear = nn.Linear (20, 30, bias=True) # input_feature: 20; output_feature: 30

# Forward computation
output_tensor = linear (input_tensor) # output_tensor shape: [128, 30]

# Print the shape and value of the output.
print("Linear output shape ", output_tensor.shape)
print(ouput_tensor)

Parent topic: Examples