AI开发平台ModelArts-PyTorch:训练模型

时间：2023-11-01 16:20:34

AI开发平台ModelArts

训练模型

from __future__ import print_functionimport argparseimport torchimport torch.nn as nnimport torch.nn.functional as Fimport torch.optim as optimfrom torchvision import datasets, transforms# 定义网络结构class Net(nn.Module):    def __init__(self):        super(Net, self).__init__()        # 输入第二维需要为784        self.hidden1 = nn.Linear(784, 5120, bias=False)        self.output = nn.Linear(5120, 10, bias=False)    def forward(self, x):        x = x.view(x.size()[0], -1)        x = F.relu((self.hidden1(x)))        x = F.dropout(x, 0.2)        x = self.output(x)        return F.log_softmax(x)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.cross_entropy(output, target)        loss.backward()        optimizer.step()        if batch_idx % 10 == 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').item()  # sum up batch loss            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability            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)))device = torch.device("cpu")batch_size=64kwargs={}train_loader = torch.utils.data.DataLoader(    datasets.MNIST('.', train=True, download=True,                   transform=transforms.Compose([                       transforms.ToTensor()                   ])),    batch_size=batch_size, shuffle=True, **kwargs)test_loader = torch.utils.data.DataLoader(    datasets.MNIST('.', train=False, transform=transforms.Compose([        transforms.ToTensor()    ])),    batch_size=1000, shuffle=True, **kwargs)model = Net().to(device)optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)optimizer = optim.Adam(model.parameters())for epoch in range(1, 2 + 1):    train(model, device, train_loader, optimizer, epoch)    test(model, device, test_loader)

上一篇：AI开发平台ModelArts-PyTorch:推理代码

下一篇：编译构建 CodeArts Build-示例2：根据项目ID查看项目下构建任务列表:操作步骤