torch conv2d 卷积底层复现验证算法实现

努力学习计算机123

已于 2022-12-05 19:38:39 修改

阅读量654

点赞数

分类专栏：深度学习文章标签：算法深度学习 python pytorch 卷积神经网络

于 2022-12-05 00:40:01 首次发布

本文链接：https://blog.csdn.net/qq_40939814/article/details/128179849

版权

Conv2d

卷积是什么?

在这里插入图片描述

这里提供torch conv2d算法的两种python的实现方式. (含dilation, groups)

使用im2col的方式实现 (更易于实现, 更快的速度)
使用循环方式实现 (更易于理解)

代码地址: https://github.com/Jintao-Huang/ml_alg/blob/main/libs/ml/_ml_alg/_nn_functional.py

以下代码含测试代码:

import torch
import torch.nn.functional as F
from torch import Tensor
from typing import Optional, Tuple


def conv2d(
    x: Tensor,
    weight: Tensor,
    bias: Optional[Tensor] = None,
    stride: Tuple[int, int] = (1, 1),
    padding: Tuple[int, int] = (0, 0),
    dilation: Tuple[int, int] = (1, 1),
    groups: int = 1
) -> Tensor:
    """faster than conv2d_2, but more memory. (recommend)
    x: [N, Cin, Hin, Win]
    weight: [Cout, Cin//G, KH, KW]. 
    bias: [Cout]
    stride: SH, SW
    padding: PH, PW
    return: [N, Cout, Hout, Wout]
    """
    Hin, Win = x.shape[2:]
    DH, DW = dilation
    G = groups
    KH, KW = weight.shape[2:]
    KH_D, KW_D = (KH - 1) * DH + 1, (KW - 1) * DW + 1
    PH, PW = padding
    SH, SW = stride
    N, Cin = x.shape[:2]
    Cout = weight.shape[0]
    # Out = (In + 2*P − ((K-1)*D+1)) // S + 1
    Hout, Wout = (Hin + 2 * PH - KH_D) // SH + 1, (Win + 2 * PW - KW_D) // SW + 1
    assert weight.shape