捕获.PNG1081×500 25.7 KB

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1=nn.Sequential(
            nn.Conv2d(
                in_channels=1,
               out_channels=16,
               kernel_size=3,
               stride=1,
               padding=1,
        ),
           nn.ReLU(),

     )
       self.conv2 = nn.Sequential(  
            nn.Conv2d(16, 32, 3, 1, 1),  
            nn.ReLU(),  # activation
       )
    self.out = nn.Linear(32 * 7 * 7, 5)

   def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)  
       output = self.out(x)
       return output, x
cnn = CNN().cuda()
print(cnn)
EPOCH=10
optimizer = optim.SGD(cnn.parameters(), lr=0.1, momentum=0.9,weight_decay=1e-4,nesterov=False)
criterion = torch.nn.CrossEntropyLoss()
def train(cnn,train_loader,EPOCH):
     cnn.train()
     for epoch in range(EPOCH):
         for step,(data,target) in enumerate(train_loader):
             data,target=data.cuda(),target.cuda()
             data,target=Variable(data),Variable(target)
             optimizer.zero_grad()
             output=cnn(data)
             loss=criterion(output,target)
             loss.backward()
             optimizer.step()
             if step % 10 == 0:
                  print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                             epoch, step * len(data), len(train_loader.dataset),
                              100. * step / len(train_loader), loss.data[0]))
def test(cnn, test_loader):
    cnn.eval()
    test_loss = 0
    correct = 0
    for data, target in test_loader:
        data, target = data.cuda(), target.cuda()
        data, target = Variable(data, volatile=True), Variable(target)
        output = cnn(data)
        test_loss += criterion(output, target).data[0]  # Variable.data
        pred = output.data.max(1)[1]  # get the index of the max log-probability
        correct += pred.eq(target.data).cpu().sum()
    test_loss = test_loss
    test_loss /= len(test_loader)  # loss function already averages over batch size
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))
if __name__ == "__main__":
    train(cnn, train_loader, EPOCH=10)
    test(cnn, test_loader)