该代码示例包含了训练、测试和验证过程的实现。首先,定义一个叫做Model的模型,它包含两个全连接层。然后定义了train函数,它遍历训练集,计算损失并更新模型参数。在同样的方式下,实现validate和test函数,它们分别计算验证集和测试集的损失和准确率。
在main函数中,定义了一系列超参数和设备,然后将训练集、验证集和测试集加载进来。接着,创建一个新的Model对象,定义损失函数和优化器,并开始训练和验证过程。最后,调用test函数进行测试。
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
# 定义模型
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
self.fc1 = nn.Linear(100, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
# 训练
def train(model, optimizer, criterion, train_loader, device):
total_loss = 0
# 将模型置为训练模式
model.train()
# 遍历训练集
for data, target in train_loader:
# 将数据移到设备上
data, target = data.to(device), target.to(device)
# 清零梯度
optimizer.zero_grad()
# 前向传播
output = model(data)
# 计算损失
loss = criterion(output, target)
# 反向传播
loss.backward()
# 更新参数
optimizer.step()
# 统计损失
total_loss += loss.item()
# 计算平均损失
avg_loss = total_loss / len(train_loader)
return avg_loss
# 验证
def validate(model, criterion, val_loader, device):
total_loss = 0
# 将模型置为验证模式
model.eval()
with torch.no_grad():
# 遍历验证集
for data, target in val_loader:
# 将数据移到设备上
data, target = data.to(device), target.to(device)
# 前向传播
output = model(data)
# 计算损失
loss = criterion(output, target)
# 统计损失
total_loss += loss.item()
# 计算平均损失
avg_loss = total_loss / len(val_loader)
return avg_loss
# 测试
def test(model, criterion, test_loader, device):
total_loss = 0
total_correct = 0
# 将模型置为测试模式
model.eval()
with torch.no_grad():
# 遍历测试集
for data, target in test_loader:
# 将数据移到设备上
data, target = data.to(device), target.to(device)
# 前向传播
output = model(data)
# 计算损失
loss = criterion(output, target)
# 统计损失
total_loss += loss.item()
# 计算准确率
_, predicted = torch.max(output.data, 1)
total_correct += (predicted == target).sum().item()
# 计算平均损失
avg_loss = total_loss / len(test_loader)
# 计算准确率
accuracy = total_correct / len(test_loader.dataset)
return avg_loss, accuracy
# 主函数
def main():
# 定义超参数
epochs = 10
lr = 0.01
batch_size = 32
# 定义设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 定义训练集、验证集、测试集
train_loader = DataLoader(train_data, batch_size=batch_size, shuffle=True)
val_loader = DataLoader(val_data, batch_size=batch_size, shuffle=False)
test_loader = DataLoader(test_data, batch_size=batch_size, shuffle=False)
# 定义模型
model = Model().to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=lr)
# 训练和验证
for epoch in range(1, epochs + 1):
train_loss = train(model, optimizer, criterion, train_loader, device)
val_loss = validate(model, criterion, val_loader, device)
print(f"Epoch: {epoch}, Train loss: {train_loss:.4f}, Val loss: {val_loss:.4f}")
# 测试
test_loss, test_accuracy = test(model, criterion, test_loader, device)
print(f"Test loss: {test_loss:.4f}, Test accuracy: {test_accuracy:.4f}")
if __name__ == '__main__':
main()
这是一个基本框架,使用这个框架,就可以完成很多使用pytorch的项目。