SPDConv代码

最新推荐文章于 2025-06-18 21:47:18 发布

weixin_45734021

最新推荐文章于 2025-06-18 21:47:18 发布

阅读量345

点赞数 5

CC 4.0 BY-SA版权

文章标签： python 开发语言计算机视觉

本文链接：https://blog.csdn.net/weixin_45734021/article/details/147712703

在这里插入图片描述

import torch
import torch.nn as nn


def autopad(k, p=None, d=1):  # kernel, padding, dilation
    # Pad to 'same' shape outputs
    if d > 1:
        k = d * (k - 1) + 1 if isinstance(k, int) else [d * (x - 1) + 1 for x in k]  # actual kernel-size
    if p is None:
        p = k // 2 if isinstance(k, int) else [x // 2 for x in k]  # auto-pad
    return p

class SPDConv(nn.Module):
    """Standard convolution with args(ch_in, ch_out, kernel, stride, padding, groups, dilation, activation)."""
    default_act = nn.SiLU()  # default activation

    def __init__(self, c1, c2, k=1, s=1, p=None, g=1, d=1, act=True):
        """Initialize Conv layer with given arguments including activation."""
        super().__init__()
        c1 = c1 * 4
        self.conv = nn.Conv2d(c1, c2, k, s, autopad(k, p, d), groups=g, dilation=d, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = self.default_act if act is True else act if isinstance(act, nn.Module) else nn.Identity()

    def forward(self, x):
        x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
        """Apply convolution, batch normalization and activation to input tensor."""
        return self.act(self.bn(self.conv(x)))

    def forward_fuse(self, x):
        """Perform transposed convolution of 2D data."""
        x = torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)
        return self.act(self.conv(x))

if __name__ == '__main__':
    # 初始化 CPAM 模型，输入通道数为 64
    model =SPDConv(64,64)
    # 定义两个输入张量，形状为 [1, 64, 32, 32]

    input1 = torch.randn(1, 64, 32, 32)
    # 将输入张量传入模型，得到输出
    output = model(input1)

    # 打印输入和输出的形状
    print("Input 1 shape:", input1.shape)

    print("Output shape:", output.shape)
    # 计算模型的总参数量，并打印
    total_params = sum(p.numel() for p in model.parameters())
    print(f'Total parameters: {total_params / 1e6:.2f}M')