Simon Shi的小站

def conv_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的卷积层和tensorflow的卷积层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """
    # 创建pytorch卷积层
    torch_conv = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
    # [kernel_number, kernel_channel, kernel_height, kernel_width]
    # 卷积层的weights
    torch_conv_weight = torch_conv.weight
    # 卷积层的bias
    torch_conv_bias = torch_conv.bias

    # 创建tensorflow卷积层
    tf_conv = tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same')
    tf_conv.build([1, 5, 5, 3])
    # 将pytorch的卷积层权重进行转换并载入tf的卷积层中
    # to [kernel_height, kernel_width, kernel_channel, kernel_number]
    value = np.transpose(torch_conv_weight.detach().numpy(), (2, 3, 1, 0)).astype(np.float32)
    tf_conv.set_weights([value, torch_conv_bias.detach().numpy()])

    # 计算pytorch卷积层的输出
    # [B, C, H, W]
    v1 = torch_conv(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)
    # [H, W, C]
    v1 = np.transpose(v1, (1, 2, 0))

    # 计算tensorflow卷积层的输出
    # [B, H, W, C]
    v2 = tf_conv(tf_image).numpy()
    # [H, W, C]
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-05)
    print("convolution layer test is great!")

def dw_conv_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的dw卷积层和tensorflow的dw卷积层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """
    # 创建pytorch的dw卷积层
    torch_conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, groups=3)
    # [kernel_number, kernel_channel, kernel_height, kernel_width]
    # dw卷积层的weights
    torch_conv_weight = torch_conv.weight
    # dw卷积层的bias
    torch_conv_bias = torch_conv.bias

    # 创建tensorflow的dw卷积层
    tf_conv = tf.keras.layers.DepthwiseConv2D(kernel_size=3, padding='same')
    tf_conv.build([1, 5, 5, 3])
    # 将pytorch的dw卷积层权重进行转换并载入tf的dw卷积层中
    # to [kernel_height, kernel_width, kernel_number, kernel_channel]
    value = np.transpose(torch_conv_weight.detach().numpy(), (2, 3, 0, 1)).astype(np.float32)
    tf_conv.set_weights([value, torch_conv_bias.detach().numpy()])

    # 计算pytorch卷积层的输出
    # [B, C, H, W]
    v1 = torch_conv(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)
    # [H, W, C]
    v1 = np.transpose(v1, (1, 2, 0))

    # 计算tensorflow卷积层的输出
    # [B, H, W, C]
    v2 = tf_conv(tf_image).numpy()
    # [H, W, C]
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-05)
    print("depthwise convolution layer test is great!")

def bn_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的bn层和tensorflow的bn层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """
    # 创建pytorch的bn层
    torch_bn = nn.BatchNorm2d(num_features=3, eps=1e-5)
    # 随机初始化bn的参数
    nn.init.uniform_(torch_bn.weight, a=1, b=5)
    nn.init.uniform_(torch_bn.bias, a=0.05, b=0.1)
    nn.init.uniform_(torch_bn.running_mean, a=0.05, b=0.1)
    nn.init.uniform_(torch_bn.running_var, a=1, b=5)
    # bn的weights
    torch_bn_weight = torch_bn.weight
    # bn的bias
    torch_bn_bias = torch_bn.bias
    # bn的running_mean
    torch_bn_mean = torch_bn.running_mean
    # bn的running_var
    torch_bn_var = torch_bn.running_var

    # 创建tensorflow的bn层
    tf_bn = tf.keras.layers.BatchNormalization(epsilon=1e-5)
    tf_bn.build([1, 5, 5, 3])
    # 将pytorch的bn权重载入tf的bn中
    tf_bn.set_weights([torch_bn_weight.detach().numpy(),
                       torch_bn_bias.detach().numpy(),
                       torch_bn_mean.detach().numpy(),
                       torch_bn_var.detach().numpy()])

    # 计算pytorch bn的输出
    # [B, C, H, W]
    torch_bn.eval()
    v1 = torch_bn(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)
    # [H, W, C]
    v1 = np.transpose(v1, (1, 2, 0))

    # 计算tensorflow bn的输出
    # [B, H, W, C]
    v2 = tf_bn(tf_image, training=False).numpy()
    # [H, W, C]
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-04)
    print("bn layer test is great!")

def fc_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的fc层和tensorflow的fc层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """

    # mean height and width dim
    torch_image = torch.mean(torch_image, dim=[2, 3])
    tf_image = np.mean(tf_image, axis=(1, 2))

    # 创建pytorch的fc卷积层
    torch_fc = nn.Linear(in_features=3, out_features=5)
    # [output_units, input_units]
    # fc层的weights
    torch_fc_weight = torch_fc.weight
    # fc层的bias
    torch_fc_bias = torch_fc.bias

    # 创建tensorflow的fc层
    tf_fc = tf.keras.layers.Dense(units=5)
    tf_fc.build([1, 3])
    # 将pytorch的fc层权重进行转换并载入tf的fc层中
    # to [input_units, output_units]
    value = np.transpose(torch_fc_weight.detach().numpy(), (1, 0)).astype(np.float32)
    tf_fc.set_weights([value, torch_fc_bias.detach().numpy()])

    # 计算pytorch fc的输出
    # [B, C]
    v1 = torch_fc(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)

    # 计算tensorflow fc的输出
    # [C, B]
    v2 = tf_fc(tf_image).numpy()
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-05)
    print("fc layer test is great!")

import tensorflow as tf
import torch
from torch import nn
import numpy as np


def conv_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的卷积层和tensorflow的卷积层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """
    # 创建pytorch卷积层
    torch_conv = nn.Conv2d(in_channels=3, out_channels=32, kernel_size=3, padding=1)
    # [kernel_number, kernel_channel, kernel_height, kernel_width]
    # 卷积层的weights
    torch_conv_weight = torch_conv.weight
    # 卷积层的bias
    torch_conv_bias = torch_conv.bias

    # 创建tensorflow卷积层
    tf_conv = tf.keras.layers.Conv2D(filters=32, kernel_size=3, padding='same')
    tf_conv.build([1, 5, 5, 3])
    # 将pytorch的卷积层权重进行转换并载入tf的卷积层中
    # to [kernel_height, kernel_width, kernel_channel, kernel_number]
    value = np.transpose(torch_conv_weight.detach().numpy(), (2, 3, 1, 0)).astype(np.float32)
    tf_conv.set_weights([value, torch_conv_bias.detach().numpy()])

    # 计算pytorch卷积层的输出
    # [B, C, H, W]
    v1 = torch_conv(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)
    # [H, W, C]
    v1 = np.transpose(v1, (1, 2, 0))

    # 计算tensorflow卷积层的输出
    # [B, H, W, C]
    v2 = tf_conv(tf_image).numpy()
    # [H, W, C]
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-05)
    print("convolution layer test is great!")


def dw_conv_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的dw卷积层和tensorflow的dw卷积层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """
    # 创建pytorch的dw卷积层
    torch_conv = nn.Conv2d(in_channels=3, out_channels=3, kernel_size=3, padding=1, groups=3)
    # [kernel_number, kernel_channel, kernel_height, kernel_width]
    # dw卷积层的weights
    torch_conv_weight = torch_conv.weight
    # dw卷积层的bias
    torch_conv_bias = torch_conv.bias

    # 创建tensorflow的dw卷积层
    tf_conv = tf.keras.layers.DepthwiseConv2D(kernel_size=3, padding='same')
    tf_conv.build([1, 5, 5, 3])
    # 将pytorch的dw卷积层权重进行转换并载入tf的dw卷积层中
    # to [kernel_height, kernel_width, kernel_number, kernel_channel]
    value = np.transpose(torch_conv_weight.detach().numpy(), (2, 3, 0, 1)).astype(np.float32)
    tf_conv.set_weights([value, torch_conv_bias.detach().numpy()])

    # 计算pytorch卷积层的输出
    # [B, C, H, W]
    v1 = torch_conv(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)
    # [H, W, C]
    v1 = np.transpose(v1, (1, 2, 0))

    # 计算tensorflow卷积层的输出
    # [B, H, W, C]
    v2 = tf_conv(tf_image).numpy()
    # [H, W, C]
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-05)
    print("depthwise convolution layer test is great!")


def bn_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的bn层和tensorflow的bn层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """
    # 创建pytorch的bn层
    torch_bn = nn.BatchNorm2d(num_features=3, eps=1e-5)
    # 随机初始化bn的参数
    nn.init.uniform_(torch_bn.weight, a=1, b=5)
    nn.init.uniform_(torch_bn.bias, a=0.05, b=0.1)
    nn.init.uniform_(torch_bn.running_mean, a=0.05, b=0.1)
    nn.init.uniform_(torch_bn.running_var, a=1, b=5)
    # bn的weights
    torch_bn_weight = torch_bn.weight
    # bn的bias
    torch_bn_bias = torch_bn.bias
    # bn的running_mean
    torch_bn_mean = torch_bn.running_mean
    # bn的running_var
    torch_bn_var = torch_bn.running_var

    # 创建tensorflow的bn层
    tf_bn = tf.keras.layers.BatchNormalization(epsilon=1e-5)
    tf_bn.build([1, 5, 5, 3])
    # 将pytorch的bn权重载入tf的bn中
    tf_bn.set_weights([torch_bn_weight.detach().numpy(),
                       torch_bn_bias.detach().numpy(),
                       torch_bn_mean.detach().numpy(),
                       torch_bn_var.detach().numpy()])

    # 计算pytorch bn的输出
    # [B, C, H, W]
    torch_bn.eval()
    v1 = torch_bn(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)
    # [H, W, C]
    v1 = np.transpose(v1, (1, 2, 0))

    # 计算tensorflow bn的输出
    # [B, H, W, C]
    v2 = tf_bn(tf_image, training=False).numpy()
    # [H, W, C]
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-04)
    print("bn layer test is great!")


def fc_test(torch_image, tf_image):
    """
    测试转换权重后pytorch的fc层和tensorflow的fc层输出是否一致
    :param torch_image:
    :param tf_image:
    :return:
    """

    # mean height and width dim
    torch_image = torch.mean(torch_image, dim=[2, 3])
    tf_image = np.mean(tf_image, axis=(1, 2))

    # 创建pytorch的fc卷积层
    torch_fc = nn.Linear(in_features=3, out_features=5)
    # [output_units, input_units]
    # fc层的weights
    torch_fc_weight = torch_fc.weight
    # fc层的bias
    torch_fc_bias = torch_fc.bias

    # 创建tensorflow的fc层
    tf_fc = tf.keras.layers.Dense(units=5)
    tf_fc.build([1, 3])
    # 将pytorch的fc层权重进行转换并载入tf的fc层中
    # to [input_units, output_units]
    value = np.transpose(torch_fc_weight.detach().numpy(), (1, 0)).astype(np.float32)
    tf_fc.set_weights([value, torch_fc_bias.detach().numpy()])

    # 计算pytorch fc的输出
    # [B, C]
    v1 = torch_fc(torch_image).detach().numpy()
    v1 = np.squeeze(v1, axis=0)

    # 计算tensorflow fc的输出
    # [C, B]
    v2 = tf_fc(tf_image).numpy()
    v2 = np.squeeze(v2, axis=0)

    # 检查pytorch和tensorflow的输出结果是否一致
    np.testing.assert_allclose(v1, v2, rtol=1e-03, atol=1e-05)
    print("fc layer test is great!")


def main():
    image = np.random.rand(5, 5, 3)
    torch_image = np.transpose(image, (2, 0, 1)).astype(np.float32)
    # [B, C, H, W]
    torch_image = torch.unsqueeze(torch.as_tensor(torch_image), dim=0)
    # [B, H, W, C]
    tf_image = np.expand_dims(image, axis=0)

    conv_test(torch_image, tf_image)
    dw_conv_test(torch_image, tf_image)
    bn_test(torch_image, tf_image)
    fc_test(torch_image, tf_image)


if __name__ == '__main__':
    main()

def update(self, state, action, reward, next_state, done):
    state = torch.FloatTensor(state).unsqueeze(0)
    next_state = torch.FloatTensor(next_state).unsqueeze(0)
    action = torch.LongTensor([action])
    reward = torch.FloatTensor([reward])
    done = torch.FloatTensor([int(done)])

    # 计算advantage
    _, next_state_value = self.actor_critic(next_state)
    _, state_value = self.actor_critic(state)
    advantage = reward + self.gamma * next_state_value * (1 - done) - state_value

    # 计算actor和critic的loss
    log_prob, _ = self.actor_critic(state)
    actor_loss = -(log_prob[0][action] * advantage).mean()
    critic_loss = advantage.pow(2).mean()
    loss = actor_loss + critic_loss

    # 更新actor和critic的参数
    self.optimizer.zero_grad()
    loss.backward()
    self.optimizer.step()