import torch
from torch import nn
'''
测试池化和卷积组合的分类层
'''
class MyNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = nn.Conv2d(3, 32, 3, padding=1)
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.classifier = nn.Linear(32, 2)
    def forward(self, x):
        x = self.conv(x)
        x = self.avg_pool(x)
        x = x.view(x.size(0), -1) # 展开所有元素
        out = self.classifier(x)
        return out
if __name__ == '__main__':
    from torchsummary import summary
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    x = torch.rand(size=(1, 3, 7, 7)).to(device)
    net = MyNet().to(device)
    summary(net, (3, 7, 7))

import torch
from torch import nn
'''
测试池化和卷积组合的分类层
'''
class MyNet(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv = nn.Conv2d(3, 32, 3, padding=1)
        self.classifier = nn.Sequential(
            nn.AdaptiveAvgPool2d(1),
            nn.Flatten(),
            nn.Linear(32, 2)
        )
    def forward(self, x):
        x = self.conv(x)
        out = self.classifier(x)
        return out
if __name__ == '__main__':
    from torchsummary import summary
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    x = torch.rand(size=(1, 3, 7, 7)).to(device)
    net = MyNet().to(device)
    out = net(x)
    summary(net, (3, 7, 7))