text: import hiddenlayer as hl
  import torch
  import torch.nn as nn
  import torchvision
  import torchvision.transforms as transforms
  import torch.nn.functional as F
  
  # Device configuration
  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  
  # Hyper parameters
  num_epochs = 5
  num_classes = 10
  batch_size = 64
  learning_rate = 0.001
  
  # 数据集
  train_dataset = torchvision.datasets.FashionMNIST(
      root='datasets/', train=True, transform=transforms.ToTensor(), download=True)
  test_dataset = torchvision.datasets.FashionMNIST(
      root='datasets/', train=False, transform=transforms.ToTensor())
  
  # Data loader
  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=False)
  
  class LeNet(nn.Module):
      def __init__(self, in_channels=3, num_classes=10):
          super(LeNet, self).__init__()
          self.conv1 = nn.Conv2d(in_channels, 6, kernel_size=5)
          self.maxp1 = nn.MaxPool2d(kernel_size=2, stride=2)
          self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
          self.fc1 = nn.Linear(256, 120)
          self.fc2 = nn.Linear(120, 84)
          self.fc3 = nn.Linear(84, 10)
  
      def forward(self, x):
          out = self.conv1(x)
          out = self.maxp1(F.sigmoid(out))
          out = self.conv2(out)
          out = self.maxp1(F.sigmoid(out))
          out = out.view(out.size(0), -1)
          out = self.fc1(out)
          out = self.fc2(F.sigmoid(out))
          out = self.fc3(F.sigmoid(out))
          return out
  
  
  model = LeNet(1, num_classes).to(device)
  h1 = hl.build_graph(model, torch.zeros(64, 1, 28, 28).to(device))
  h1.save('images/lenet.png', format='png')
  
  # Loss and optimizer
  criterion = nn.CrossEntropyLoss()
  optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
  
  # Train the model
  total_step = len(train_loader)
  for epoch in range(num_epochs):
      for i, (images, labels) in enumerate(train_loader):
          images = images.to(device)
          labels = labels.to(device)
  
          # Forward pass
          outputs = model(images)
          loss = criterion(outputs, labels)
  
          # Backward and optimize
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()
  
          if (i+1) % 100 == 0:
              print('Epoch [{}/{}], Step [{}/{}], Loss: {:.4f}'
                    .format(epoch+1, num_epochs, i+1, total_step, loss.item()))
  
  # Test the model
  model.eval()
  with torch.no_grad():
      correct = 0
      for images, labels in test_loader:
          images = images.to(device)
          labels = labels.to(device)
          outputs = model(images)
          _, predicted = torch.max(outputs.data, 1)
          correct += (predicted == labels).sum().item()
  
      print('Test Accuracy of the model is: {} %'.format(100 * correct / len(test_loader.dataset)))
  