text: import torch
  import os
  import random
  from torch import nn
  from options import args_parser
  from torchvision import datasets, transforms
  from torch.utils.data import DataLoader
  
  ## 初始化一些参数
  args = args_parser()
  random.seed(args.seed)
  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  print(device)
  
  ## 分类器
  class Classifier (nn.Module):
      def __init__(self, nc,n_class):
          super(Classifier, self).__init__()
  
          self.main = nn.Sequential(
              nn.Conv2d(nc,64,4,2,1),
              nn.BatchNorm2d(64),
              # nn.LeakyReLU(0.2,inplace=False),
              nn.ReLU(),
  
              nn.Conv2d(64,128,4,2,1),
              nn.BatchNorm2d(128),
              # nn.LeakyReLU(0.2,inplace=False),
              nn.ReLU(),
  
              nn.Conv2d(128,256,4,2,1),
              nn.BatchNorm2d(256),
              # nn.LeakyReLU(0.2,inplace=False),
              nn.ReLU(),
          )
          self.C = nn.Sequential(
              nn.Linear(3 * 3 * 256, n_class)
          )
  
      def forward(self, x):
          x = self.main(x).view(x.shape[0], -1)
          x1 = self.C(x)
          return x1
  if __name__ == '__main__':
      train_data = datasets.MNIST(root='./mnist',train=True,transform=transforms.Compose([transforms.ToTensor()]))
      train_loader = DataLoader(train_data, batch_size=args.local_bs, shuffle=True)
      C = Classifier(args.cn,args.num_class).to(device)
      optimizerC = torch.optim.Adam(C.parameters(), lr=args.local_lr)
      aux_criterion = nn.CrossEntropyLoss()
      for epoch in range(args.epochs):
          total = 0
          correct =0
          for step, (x, y) in enumerate(train_loader):
              x = x.to(device)
              y = y.to(device)
              C.train()
  
              C_x = C(x)
              C_loss = aux_criterion(C_x,y)
  
              optimizerC.zero_grad()
              C_loss.backward()
              optimizerC.step()
              
              ## 观察在训练集上准确率的变化
              C_x = C(x)
              _,pred = C_x.max(1)
              total += y.size(0)
              correct += (pred == y).sum().item()
              # ====================================================================
          print("epoch:%d------accuracy:%.2f"%(epoch,(100.0 * correct / total)))
          save_dirs = 'local/simc/'
          if not os.path.exists(save_dirs):
              os.makedirs(save_dirs)
          torch.save(C, save_dirs+'C_epoch_%d.pth' %epoch)