代码语言: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)