测试caffe和pytorch结果的一致性

测试caffe和pytorch结果的一致性

以最简单的卷积为例。

caffe模型

name: "conv_test"
layer {
  name: "data"
  type: "Input"
  top: "data"
  input_param { shape: { dim: 1 dim: 32 dim: 21 dim: 21 } }
}
layer {
  name: "conv1"
  type: "Convolution"
  bottom: "data"
  top: "conv1"
  convolution_param {
    num_output: 64
    kernel_size: 3
    stride: 1
  }
}

pytorch模型和一致性测试代码

import torch
import caffe
import numpy as np


class TorchConv(torch.nn.Module):
    def __init__(self):
        super(TorchConv, self).__init__()
        self.layer1 = torch.nn.Conv2d(32, 64, (3, 3))

    def forward(self, x):
        out = self.layer1(x)
        return out


def torch_conv_layer_test():
    # 初始化网络
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    net = TorchConv().to(device).eval()

    # 初始化随机种子
    np_weight = np.random.seed(1)

    # 权重
    weight = net.layer1.weight
    bias = net.layer1.bias
    np_weight = np.random.randn(*weight.shape)
    net.layer1.weight = torch.nn.Parameter(torch.from_numpy(np_weight).to(device))
    np_bias = np.random.randn(*bias.shape)
    net.layer1.bias = torch.nn.Parameter(torch.from_numpy(np_bias).to(device))

    # 输入
    input_tonsor = torch.from_numpy(np.random.randn(1, 32, 21, 21)).to(device)

    # 输出
    output = net(input_tonsor)
    output_np = output.detach().cpu().numpy()

    return output_np


def caffe_conv_layer_test():
    # 初始化网络
    net = caffe.Net('consistency/conv.prototxt', caffe.TEST)

    # 初始化随机种子
    np_weight = np.random.seed(1)

    # 权重
    conv_layer_params = list(net.params.values())[0]
    weight = conv_layer_params[0].data
    weight[...] = np.random.randn(*weight.shape)
    bias = conv_layer_params[1].data
    bias[...] = np.random.randn(*bias.shape)

    # 输入
    input_data = net.blobs[net.inputs[0]].data
    input_data[...] = np.random.randn(1, 32, 21, 21)

    # 输出
    output = net.forward()
    output_np = np.array(output[net.outputs[0]])

    return output_np


def main():
    caffe_out = caffe_conv_layer_test()
    torch_out = torch_conv_layer_test()
    print()

    # 计算两个输出之间的均方误差
    mse = np.square(torch_out - caffe_out).mean(axis=None)
    print("MSE is", mse)


if __name__ == "__main__":
    main()

使用元素的均方误差来衡量差异度，最终的打印结果为MSE is 2.7442352745900254e-11，差异很小。