下面是一个示例代码,演示了如何编写一个具有模型拟合和预测功能的PyTorch神经网络类。
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
class NeuralNet(torch.nn.Module):
def __init__(self):
super(NeuralNet, self).__init__()
self.fc1 = torch.nn.Linear(10, 20)
self.fc2 = torch.nn.Linear(20, 1)
def forward(self, x):
x = F.relu(self.fc1(x))
x = self.fc2(x)
return x
def fit(self, train_loader, optimizer, criterion, epochs):
self.train()
for epoch in range(epochs):
running_loss = 0.0
for i, data in enumerate(train_loader, 0):
inputs, labels = data
optimizer.zero_grad()
outputs = self(inputs.float())
loss = criterion(outputs, labels.float())
loss.backward()
optimizer.step()
running_loss += loss.item()
print('Epoch %d: loss=%.3f' % (epoch+1, running_loss/len(train_loader)))
def predict(self, test_loader):
self.eval()
predictions = []
with torch.no_grad():
for i, data in enumerate(test_loader, 0):
inputs, _ = data
outputs = self(inputs.float())
predictions.append(outputs.numpy())
return predictions
在上面的代码中,我们首先定义了一个名为NeuralNet的PyTorch神经网络类。该类初始化了两个线性层并定义了前向传播函数。然后,我们有fit函数,其目的是用训练数据来拟合模型。该函数遍历训练数据集中的所有epoch,并对每个iteration执行正向和反向传播。最后,我们有predict函数,用于在训练完成后对测试数据集进行预测。
我们可以将上述类用于我们自己的数据集训