content: import torch
  class ConvNet(nn.Module):
      """Small ConvNet for MNIST."""
  
      def __init__(self):
          super(ConvNet, self).__init__()
          self.conv1 = nn.Conv2d(1, 3, kernel_size=3)
          self.fc = nn.Linear(192, 10)
  
      def forward(self, x):
          x = F.relu(F.max_pool2d(self.conv1(x), 3))
          x = x.view(-1, 192)
          x = self.fc(x)
          return F.log_softmax(x, dim=1)
  
      def get_weights(self):
          return {k: v.cpu() for k, v in self.state_dict().items()}
  
      def set_weights(self, weights):
          self.load_state_dict(weights)
  
      def get_gradients(self):
          grads = []
          for p in self.parameters():
              grad = None if p.grad is None else p.grad.data.cpu().numpy()
              grads.append(grad)
          return grads
  
      def set_gradients(self, gradients):
          for g, p in zip(gradients, self.parameters()):
              if g is not None:
                  p.grad = torch.from_numpy(g)
  
  @ray.remote
  class ParameterServer(object):
      def __init__(self, lr):
          self.model = ConvNet()
          self.optimizer = torch.optim.SGD(self.model.parameters(), lr=lr)
  
      def apply_gradients(self, *gradients):
          summed_gradients = [
              np.stack(gradient_zip).sum(axis=0)
              for gradient_zip in zip(*gradients)
          ]
          self.optimizer.zero_grad()
          self.model.set_gradients(summed_gradients)
          self.optimizer.step()
          return self.model.get_weights()
  
      def get_weights(self):
          return self.model.get_weights()
  
  @ray.remote
  class DataWorker(object):
      def __init__(self):
          self.model = ConvNet()
          self.data_iterator = iter(get_data_loader()[0])
  
      def compute_gradients(self, weights):
          self.model.set_weights(weights)
          try:
              data, target = next(self.data_iterator)
          except StopIteration:  # When the epoch ends, start a new epoch.
              self.data_iterator = iter(get_data_loader()[0])
              data, target = next(self.data_iterator)
          self.model.zero_grad()
          output = self.model(data)
          loss = F.nll_loss(output, target)
          loss.backward()
          return self.model.get_gradients()
  
  # 用户Driver程序，同步PS训练
  print("Running Synchronous Parameter Server Training.")
  iterations = 200
  num_workers = 2
  
  ray.init()
  ps = ParameterServer.remote(1e-2)
  workers = [DataWorker.remote() for i in range(num_workers)]
  model = ConvNet()
  test_loader = get_data_loader()[1]
  
  
  current_weights = ps.get_weights.remote()
  for i in range(iterations):
      gradients = [
          worker.compute_gradients.remote(current_weights) for worker in workers
      ]
      # Calculate update after all gradients are available.
      current_weights = ps.apply_gradients.remote(*gradients)
  
      if i % 10 == 0:
          # Evaluate the current model.
          model.set_weights(ray.get(current_weights))
          accuracy = evaluate(model, test_loader)
          print("Iter {}: \taccuracy is {:.1f}".format(i, accuracy))
  
  print("Final accuracy is {:.1f}.".format(accuracy))
  # Clean up Ray resources and processes before the next example.
  ray.shutdown()
  
  
  # 用户Driver程序，异步PS训练
  print("Running Asynchronous Parameter Server Training.")
  
  ray.init()
  ps = ParameterServer.remote(1e-2)
  workers = [DataWorker.remote() for i in range(num_workers)]
  
  current_weights = ps.get_weights.remote()
  
  gradients = {}
  for worker in workers:
      gradients[worker.compute_gradients.remote(current_weights)] = worker
  
  for i in range(iterations * num_workers):
      ready_gradient_list, _ = ray.wait(list(gradients))
      ready_gradient_id = ready_gradient_list[0]  # 按照任务完成的顺序
      worker = gradients.pop(ready_gradient_id)
  
      # Compute and apply gradients.
      current_weights = ps.apply_gradients.remote(*[ready_gradient_id])
      gradients[worker.compute_gradients.remote(current_weights)] = worker
  
      if i % 10 == 0:
          # Evaluate the current model after every 10 updates.
          model.set_weights(ray.get(current_weights))
          accuracy = evaluate(model, test_loader)
          print("Iter {}: \taccuracy is {:.1f}".format(i, accuracy))
  
  print("Final accuracy is {:.1f}.".format(accuracy))