定义网络
import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
def __init__(self):
'''
定义网络
'''
super(Net, self).__init__()
# 卷积层
# self.conv = nn.Conv2d(input_channel, output_channel, kernel)
# 池化层
# self.pool = nn.MaxPool2d(kernel_1, kernel_2)
# 线性层
# self.linear = nn.Linear(input_dim, output_dim)
def forward(self, x):
'''
定义前向传播
'''
#x = F.relu(self.conv(x))
#x = self.pool(x)
#x = self.linear(x)
return x
net = Net()
训练网络
import torch.optim as optim
# 定义损失函数
criterion = nn.CrossEntropyLoss()
# 定义优化器
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)
# n次epoch训练
for epoch in range(n):
for i, data in enumerate(trainloader, 0):
# 用生成器trainloader实现SGD数据集
inputs, labels = data
# 使梯度为0
optimizer.zero_grad()
# forward + backward + optimize
outputs = net(inputs)
loss = criterion(outputs, labels) # forward
loss.backward() # backward
optimizer.step() # optimize
使用GPU
import torch
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 网络生成一个副本并部署在device中
net.to(device)
# 数据生成一个副本并部署在device中
data.to(device)