Сравнение Conv2D с заполнением между Tensorflow и PyTorch

Question

Я пытаюсь импортировать веса, сохраненные из модели Tensorflow, в PyTorch.До сих пор результаты были очень похожи.Я столкнулся с загадкой, когда модель требует conv2d с stride=2.

Чтобы проверить несоответствие, я настроил очень простое сравнение между TF и ​​PyTorch.Сначала я сравниваю conv2d с stride=1.

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


np.random.seed(0)
sess = tf.Session()

# Create random weights and input
weights = torch.empty(3, 3, 3, 8)
torch.nn.init.constant_(weights, 5e-2)
x = np.random.randn(1, 3, 10, 10)

weights_tf = tf.convert_to_tensor(weights.numpy(), dtype=tf.float32)
# PyTorch adopts [outputC, inputC, kH, kW]
weights_torch = torch.Tensor(weights.permute((3, 2, 0, 1)))

# Tensorflow defaults to NHWC
x_tf = tf.convert_to_tensor(x.transpose((0, 2, 3, 1)), dtype=tf.float32)
x_torch = torch.Tensor(x)

# TF Conv2D
tf_conv2d = tf.nn.conv2d(x_tf,
                         weights_tf,
                         strides=[1, 1, 1, 1],
                         padding="SAME")

# PyTorch Conv2D
torch_conv2d = F.conv2d(x_torch, weights_torch, padding=1, stride=1)

sess.run(tf.global_variables_initializer())
tf_result = sess.run(tf_conv2d)

diff = np.mean(np.abs(tf_result.transpose((0, 3, 1, 2)) - torch_conv2d.detach().numpy()))
print('Mean of Abs Diff: {0}'.format(diff))

Результат этого выполнения:

Mean of Abs Diff: 2.0443112092038973e-08

Когда я изменяю stride на 2, результаты начинаются сварьируется.

# TF Conv2D
tf_conv2d = tf.nn.conv2d(x_tf,
                         weights_tf,
                         strides=[1, 2, 2, 1],
                         padding="SAME")

# PyTorch Conv2D
torch_conv2d = F.conv2d(x_torch, weights_torch, padding=1, stride=2)

Результат этого выполнения:

Mean of Abs Diff: 0.2104552686214447

Согласно документации PyTorch, conv2d использует заполнение нулями , определенное padding аргументТаким образом, нули добавляются слева, сверху, справа и снизу от ввода в моем примере.

Если PyTorch просто добавляет заполнение с обеих сторон на основе входного параметра, его легко скопировать в Tensorflow.

# Manually add padding - consistent with PyTorch
paddings = tf.constant([[0, 0], [1, 1], [1, 1], [0, 0]])
x_tf = tf.convert_to_tensor(x.transpose((0, 2, 3, 1)), dtype=tf.float32)
x_tf = tf.pad(x_tf, paddings, "CONSTANT")

# TF Conv2D
tf_conv2d = tf.nn.conv2d(x_tf,
                         weights_tf,
                         strides=[1, 2, 2, 1],
                         padding="VALID")

Результат этого сравнения:

Mean of Abs Diff: 1.6035047067930464e-08

Это говорит мне о том, что если я каким-то образом смогу воспроизвести поведение заполнения по умолчанию из Tensorflow в PyTorch, то мои результаты будут аналогичными.

Этот вопрос проверял поведениепрокладки в Tensorflow. Документация TF объясняет, как добавляется заполнение для "ЖЕ" сверток. Я обнаружил эти ссылки при написании этого вопроса.

Теперь, когда я знаю стратегию заполнения Tensorflow, я могу реализовать ее в PyTorch.

Answer 1

Чтобы повторить поведение, размеры заполнения рассчитываются, как описано в документации Tensorflow.Здесь я проверяю поведение заполнения, устанавливая stride=2 и дополняя ввод PyTorch.

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


np.random.seed(0)
sess = tf.Session()

# Create random weights and input
weights = torch.empty(3, 3, 3, 8)
torch.nn.init.constant_(weights, 5e-2)
x = np.random.randn(1, 3, 10, 10)

weights_tf = tf.convert_to_tensor(weights.numpy(), dtype=tf.float32)
weights_torch = torch.Tensor(weights.permute((3, 2, 0, 1)))

# Tensorflow padding behavior. Assuming that kH == kW to keep this simple.
stride = 2
if x.shape[2] % stride == 0:
    pad = max(weights.shape[0] - stride, 0)
else:
    pad = max(weights.shape[0] - (x.shape[2] % stride), 0)

if pad % 2 == 0:
    pad_val = pad // 2
    padding = (pad_val, pad_val, pad_val, pad_val)
else:
    pad_val_start = pad // 2
    pad_val_end = pad - pad_val_start
    padding = (pad_val_start, pad_val_end, pad_val_start, pad_val_end)

x_tf = tf.convert_to_tensor(x.transpose((0, 2, 3, 1)), dtype=tf.float32)
x_torch = torch.Tensor(x)
x_torch = F.pad(x_torch, padding, "constant", 0)

# TF Conv2D
tf_conv2d = tf.nn.conv2d(x_tf,
                         weights_tf,
                         strides=[1, stride, stride, 1],
                         padding="SAME")

# PyTorch Conv2D
torch_conv2d = F.conv2d(x_torch, weights_torch, padding=0, stride=stride)

sess.run(tf.global_variables_initializer())
tf_result = sess.run(tf_conv2d)

diff = np.mean(np.abs(tf_result.transpose((0, 3, 1, 2)) - torch_conv2d.detach().numpy()))
print('Mean of Abs Diff: {0}'.format(diff))

Вывод:

Mean of Abs Diff: 2.2477470551507395e-08

Я не совсем уверен, почему это происходиткогда я начал писать этот вопрос, но немного прочтения прояснил это очень быстро.Я надеюсь, что этот пример может помочь другим.