Pytorch实现分组卷积

最新推荐文章于 2026-03-01 17:30:53 发布

原创

最新推荐文章于 2026-03-01 17:30:53 发布 · 945 阅读

标签

#pytorch #深度学习 #机器学习

收录于

这段代码在CUDA设备上实现了不同组数的Group Convolution，并通过torchsummary库展示了网络结构。通过time模块记录了不同组数下运算的时间，探讨了组卷积效率与组数的关系。

import torch
import torch.nn as nn
from torchsummary import summary
import time

device = torch.device('cuda')

class Group_Conv(nn.Module):
    def __init__(self, in_chs, out_chs, num_groups):
        super(Group_Conv, self).__init__()
        self.conv = nn.Conv2d(
            in_ch