最简单的包含dropout的网络

2023-11-14 11:53:18 浏览数 (1)

代码语言:javascript复制
import torch
import torch.nn as nn
import torch.optim as optim

# 定义一个简单的神经网络
class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        self.fc1 = nn.Linear(10, 5)  # 输入维度为10,输出维度为5
        self.dropout = nn.Dropout(p=0.5)  # dropout率为0.5
        self.fc2 = nn.Linear(5, 2)  # 输入维度为5,输出维度为2

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 定义一个简单的训练函数
def train(model, data, labels, optimizer, criterion, epochs):
    for epoch in range(epochs):
        optimizer.zero_grad()
        outputs = model(data)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        print(f'Epoch {epoch   1}/{epochs}, Loss: {loss.item()}')

# 生成一些虚拟数据
data = torch.randn((100, 10))  # 100个样本,每个样本有10个特征
labels = torch.randint(0, 2, (100, 2)).float()  # 100个样本,每个样本有2个类别

# 创建模型、优化器和损失函数
model = SimpleNN()
optimizer = optim.SGD(model.parameters(), lr=0.01)
criterion = nn.CrossEntropyLoss()

# 训练模型
train(model, data, labels.argmax(dim=1), optimizer, criterion, epochs=10)

0 人点赞