pytorch简单示例

中文文档

简单三层全连接网络

import torch # N is batch size; D_in is input dimension; # H is hidden dimension; D_out is output dimension. N, D_in, H, D_out = 64, 1000, 100, 10 # Create random Tensors to hold inputs and outputs. x = torch.randn(N, D_in) y = torch.randn(N, D_out) # Use the nn package to define our model and loss function. model = torch.nn.Sequential( torch.nn.Linear(D_in, H), torch.nn.ReLU(), torch.nn.Linear(H, D_out), ) loss_fn = torch.nn.MSELoss(size_average=False) # Use the optim package to define an Optimizer that will update the weights of # the model for us. Here we will use Adam; the optim package contains many other # optimization algoriths. The first argument to the Adam constructor tells the # optimizer which Tensors it should update. learning_rate = 1e-4 optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) for t in range(500): # Forward pass: compute predicted y by passing x to the model. y_pred = model(x) # Compute and print loss. loss = loss_fn(y_pred, y) print(t, loss.item()) # Before the backward pass, use the optimizer object to zero all of the # gradients for the Tensors it will update (which are the learnable weights # of the model) optimizer.zero_grad() # Backward pass: compute gradient of the loss with respect to model parameters loss.backward() # Calling the step function on an Optimizer makes an update to its parameters optimizer.step()

MNIST识别数字

import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchvision import datasets, transforms #from torch.autograd import Variable # Training settings BATCH_SIZE = 512 # 大概需要2G的显存 EPOCHS = 20 # 总共训练批次 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # MNIST Dataset train_dataset = datasets.MNIST(root='./data/', train=True, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1037,), (0.3081,)) ])) test_dataset = datasets.MNIST(root='./data/', train=False,download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1037,), (0.3081,)) ])) # Data Loader (Input Pipeline) train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=BATCH_SIZE, shuffle=True) test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=BATCH_SIZE, shuffle=True) # 定义模型 class ConvNet(nn.Module): def __init__(self): super().__init__() #1*1*28*28 self.conv1 = nn.Conv2d(1, 10, 5) self.conv2 = nn.Conv2d(10, 20, 3) self.fc1 = nn.Linear(20 * 10 * 10, 500) self.fc2 = nn.Linear(500, 10) def forward(self, x): in_size = x.size(0) out= self.conv1(x) # 1* 10 * 24 *24 out = F.relu(out) out = F.max_pool2d(out, 2, 2) # 1* 10 * 12 * 12 out = self.conv2(out) # 1* 20 * 10 * 10 out = F.relu(out) out = out.view(in_size, -1) # 1 * 2000 out = self.fc1(out) # 1 * 500 out = F.relu(out) out = self.fc2(out) # 1 * 10 out = F.log_softmax(out, dim = 1) return out # 定义训练函数 def train(model, device, train_loader, optimizer, epoch): model.train() for batch_idx, (data, target) in enumerate(train_loader): data, target = data.to(device), target.to(device) optimizer.zero_grad() output = model(data) loss = F.nll_loss(output, target) loss.backward() optimizer.step() if (batch_idx + 1) % 30 == 0: print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( epoch, batch_idx * len(data), len(train_loader.dataset), 100. * batch_idx / len(train_loader), loss.item())) # 定义测试函数 def test(model, device, test_loader): model.eval() test_loss =0 correct = 0 with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data) test_loss += F.nll_loss(output, target, reduction = 'sum') # 将一批的损失相加 pred = output.max(1, keepdim = True)[1] # 找到概率最大的下标 correct += pred.eq(target.view_as(pred)).sum().item() test_loss /= len(test_loader.dataset) print("\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%) \n".format( test_loss, correct, len(test_loader.dataset), 100.* correct / len(test_loader.dataset))) #生成模型和优化器 model = ConvNet().to(DEVICE) optimizer = optim.Adam(model.parameters()) # 最后开始训练和测试 for epoch in range(1, EPOCHS + 1): train(model, DEVICE, train_loader, optimizer, epoch) test(model, DEVICE, test_loader)

tensor

# 创建tensor torch.FloatTensor # CPU torch.cuda.FloatTensor # GPU x = torch.Tensor(3, 4) # 默认为torch.FloatTensor x = torch.FloatTensor([[1, 2, 3], [4, 5, 6]]) x = torch.IntTensor(2, 4).zero_() # 会改变tensor的函数操作会用一个下划线后缀来标示 x = torch.rand(3, 4) x = torch.randn(3, 4) x + y torch.add(x, y, out = result) y.add_(x) # numpy a = torch.ones(5) b = a.numpy() a = np.ones(5) b = torch.from_numpy(a) # 复制到gpu x = x.cuda() y = y.cuda() x + y

nn.Module

  所有网络的基类,自己定义的模型也应该继承这个类,Modules也可以包含其它Modules,允许使用树结构嵌入他们,可以将子模块赋值给模型属性.

import torch.nn as nn import torch.nn.functional as F class Model(nn.Module): def __init__(self): super(Model, self).__init__() self.conv1 = nn.Conv2d(1, 20, 5)# submodule: Conv2d self.conv2 = nn.Conv2d(20, 20, 5) def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x))