用pythorch写一段卷积神经网络代码

在PyTorch中，构建一个卷积神经网络（Convolutional Neural Network，CNN）的基本步骤通常包括定义模型、初始化权重、前向传播以及训练过程。这里是一个简单的例子，展示如何创建一个包含两个卷积层和一个全连接层的CNN：

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self, in_channels, num_classes):
        super(SimpleCNN, self).__init__()
        
        # 卷积层
        self.conv1 = nn.Conv2d(in_channels, 16, kernel_size=3, padding=1)  # 输出通道数16
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)

        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)

        # 全连接层
        self.fc = nn.Linear(32 * (in_channels // 4) * (in_channels // 4), num_classes)  # 计算池化后的维度

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)

        x = self.conv2(x)
        x = self.relu2(x)
        x = self.pool2(x)

        # 展平并传递到全连接层
        x = x.view(-1, 32 * (x.shape[2] // 4) * (x.shape[3] // 4))
        x = self.fc(x)
        
        return x

# 使用示例
model = SimpleCNN(3, 10)  # 假设输入图片有3个通道，分类10种
input_tensor = torch.randn(1, 3, 32, 32)  # 假设大小为1x3x32x32
output = model(input_tensor)

在这个例子中，`__init__`函数定义了网络结构，`forward`函数描述了数据流。注意，这只是一个基本框架，实际应用中可能还需要添加批次归一化、dropout等组件，并进行优化和损失计算。