对比tensorflow，从0开始学pytorch（三）--自定义层

瀚岳-诸葛弩2025-12-13 9:24

上文虽然实现了GMS层的效果，但是前端代码太多，太ugly，也不好复用。今天抽空看了下pytorch中怎么自定义层，很简单，比tensorflow好用。

任意文件夹创建个文件，和所有编程语言一样

一样，集成nn.Module，然后自定义一个形参

这里需要花时间搞明白torch.nn.functional下的函数和torch.nn下的类的区别，一开始有点懵，想着为什么不做高级语言当中的静态函数，想明白了也就简单了。

图中的SPP_Sizes做了类型定义，python中一般情况不需要定义类型，但不定义在后面循环就会报错，看了下pytorch自带conv2d的源码，发现源码中非常严谨，每一个变量都定义了类型。

调用就非常简单了，上一篇笔记中的冗长的代码，就可以一行调用

简化后，代码看过去顺眼多了。附上GMS封装后的源码和训练结果

python 复制代码

import torch
import torch.nn as nn
import torch.nn.functional as F


class GMS(nn.Module):
    def __init__(self, Spp_Sizes:[]):
        super().__init__()
        if len(Spp_Sizes) == 0:
            self.SPP_Sizes = [2, 3, 4]
        else:
            self.SPP_Sizes = Spp_Sizes

    def forward(self, x):
        x_gap = F.adaptive_avg_pool2d(x, (1, 1))
        x_gap = torch.flatten(x_gap, 1)

        x_gmp = F.adaptive_max_pool2d(x, (1, 1))
        x_gmp = torch.flatten(x_gmp, 1)

        x_gms = torch.cat((x_gap, x_gmp), dim=1)

        for spp_size in self.SPP_Sizes:
            x_spp = F.adaptive_max_pool2d(x, (spp_size,spp_size))
            x_spp = torch.flatten(x_spp, 1)
            x_gms = torch.cat((x_gms, x_spp), dim=1)
        return x_gms