Keras/TensorFlow 相當於 PyTorch Conv1d

Question

我目前正在將 PyTorch 代碼轉換為 TensorFlow (Keras)。 使用的層之一是 Conv1d，關於如何在 PyTorch 中使用它的描述如下

torch.nn.Conv1d(in_channels: int, out_channels: int, kernel_size: Union[T, Tuple[T]], stride: Union[T, Tuple[T]] = 1, padding: Union[T, Tuple[T]] = 0, dilation: Union[T, Tuple[T]] = 1, groups: int = 1, bias: bool = True, padding_mode: str = 'zeros')

在 Keras (TF1.15) 中，描述如下

tf.keras.layers.Conv1D(filters, kernel_size, strides=1, padding='valid', data_format='channels_last', dilation_rate=1, activation=None, use_bias=True, kernel_initializer='glorot_uniform', bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None, **kwargs)

我無法在 TensorFlow 中重現在 PyTorch 中獲得的相同輸出。 即對於 PyTorch 中的示例代碼

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

B, K, L, N = 4, 10, 128, 64
mixture = torch.randint(3, (B, K, L), dtype=torch.float32)
# L2 Norm along L axis
EPS = 1e-8
norm_coef = torch.norm(mixture, p=2, dim=2, keepdim=True)  # B x K x 1
norm_mixture = mixture / (norm_coef + EPS)  # B x K x L
# 1-D gated conv
norm_mixture = torch.unsqueeze(norm_mixture.view(-1, L), 2)  # B*K x L x 1
conv1d_U = nn.Conv1d(L, N, kernel_size=1, stride=1, bias=False)
conv_out = conv1d_U(norm_mixture)
conv_out = F.relu(conv_out)  # B*K x N x 1
mixture_w = conv_out.view(B, K, N)  # B x K x N

weights = conv1d_U.weight.data

在 TensorFlow 中，為了獲得類似的輸出尺寸，可以在下面找到代碼

import tensorflow as tf
import numpy as np

def build_net():
    # Encoder
    mixture = tf.keras.layers.Input(shape=(10, 128), name='mixture', batch_size=4)  # [B,K,L]
    norm_coef = tf.keras.backend.sqrt(tf.keras.backend.sum(mixture ** 2, axis=2, keepdims=True) + 1e-8)  # [B,K,1]
    norm_mixture = mixture / norm_coef  # [B, K, L]
    norm_mixture = tf.keras.backend.expand_dims(tf.keras.backend.reshape(norm_mixture, [-1, 128]), axis=2)  # [B*K,L,1]
    conv = tf.keras.layers.Conv1D(filters=64, kernel_size=1, activation='relu', use_bias=False, name='conv')(norm_mixture)  # [B*K,N,1]
    mixture_w = tf.keras.backend.reshape(conv, [4, -1, 64])  # [B, K, N]
    return tf.keras.models.Model(inputs=mixture, outputs=mixture_w)

model = build_net()
weights = model.get_weights()
inp = np.random.randn(4, 10, 128)
out = model.predict(inp)

比較兩種情況下的權重維度，Conv1d 操作與 TensorFlow (Keras) 明顯不同，應該如何更改 TF 代碼以反映相同的操作？

Answer 1

import tensorflow as tf
import numpy as np

def build_net():
    # Encoder
    mixture = tf.keras.layers.Input(shape=(10, 128), name='mixture', batch_size=4)  # [B,K,L]
    norm_coef = tf.keras.backend.sqrt(tf.keras.backend.sum(mixture ** 2, axis=2, keepdims=True) + 1e-8)  # [B,K,1]
    norm_mixture = mixture / norm_coef  # [B, K, L]
    norm_mixture = tf.keras.backend.expand_dims(tf.keras.backend.reshape(norm_mixture, [-1, 128]), axis=1)  # [B*K,1,L]
    conv = tf.keras.layers.Conv1D(filters=64, kernel_size=1, activation='relu', use_bias=False, name='conv')(norm_mixture)  # [B*K,1,N]
    mixture_w = tf.keras.backend.reshape(conv, [4, -1, 64])  # [B, K, N]
    return tf.keras.models.Model(inputs=mixture, outputs=mixture_w)

model = build_net()
print(model.summary())
weights = model.get_weights()
inp = np.random.randn(4, 10, 128)
out = model.predict(inp)