wehub-resource-sync
35 lines
987 B
Python
35 lines
987 B
Python
import torch
|
|
|
|
from keras.src import layers
|
|
from keras.src import testing
|
|
from keras.src.backend.common import KerasVariable
|
|
|
|
|
|
class Net(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.fc1 = layers.Dense(1)
|
|
|
|
def forward(self, x):
|
|
x = self.fc1(x)
|
|
return x
|
|
|
|
|
|
class TorchWorkflowTest(testing.TestCase):
|
|
def test_keras_layer_in_nn_module(self):
|
|
net = Net()
|
|
|
|
# Test using Keras layer in a nn.Module.
|
|
# Test forward pass
|
|
self.assertEqual(net(torch.empty(100, 10)).shape, (100, 1))
|
|
# Test KerasVariables are added as nn.Parameter.
|
|
self.assertLen(list(net.parameters()), 2)
|
|
|
|
# Test using KerasVariable as a torch tensor for torch ops.
|
|
kernel = net.fc1.kernel
|
|
transposed_kernel = torch.transpose(kernel, 0, 1)
|
|
self.assertIsInstance(kernel, KerasVariable)
|
|
self.assertIsInstance(
|
|
torch.mul(kernel, transposed_kernel), torch.Tensor
|
|
)
|