text: import os
  os.environ['OPENBLAS_NUM_THREADS'] = '1'
  
  import argparse
  import random
  
  import numpy as np
  from tensorboardX import SummaryWriter
  
  import jittor as jt
  import jittor.nn as nn
  from jittor.optim import Adam, SGD
  from jittor.lr_scheduler import MultiStepLR
  jt.flags.use_cuda = 1
  
  from tqdm import tqdm
  
  from subdivnet.dataset import SegmentationDataset
  from subdivnet.deeplab import MeshDeepLab
  from subdivnet.deeplab import MeshVanillaUnet
  from subdivnet.utils import to_mesh_tensor
  from subdivnet.utils import save_results
  from subdivnet.utils import update_label_accuracy
  from subdivnet.utils import compute_original_accuracy
  from subdivnet.utils import SegmentationMajorityVoting
  
  np.random.seed = 0
  random.seed = 0
  
  
  def train(net, optim, dataset, writer, epoch):
      net.train()
      acc = 0
      for meshes, labels, _ in tqdm(dataset, desc=str(epoch)):
          mesh_tensor = to_mesh_tensor(meshes)
          mesh_labels = jt.int32(labels)
          outputs = net(mesh_tensor)
          loss = nn.cross_entropy_loss(outputs.unsqueeze(dim=-1), mesh_labels.unsqueeze(dim=-1), ignore_index=-1)
          optim.step(loss)
  
          preds = np.argmax(outputs.data, axis=1)
          acc += np.sum((labels == preds).sum(axis=1) / meshes['Fs'])
          writer.add_scalar('loss', loss.data[0], global_step=train.step)
          train.step += 1
      acc /= dataset.total_len
      print('train acc = ', acc)
      writer.add_scalar('train-acc', acc, global_step=epoch)
  
  
  @jt.single_process_scope()
  def test(net, dataset, writer, epoch, args):
      net.eval()
      acc = 0
      oacc = 0
      label_acc = np.zeros(args.parts)
      name = args.name
      voted = SegmentationMajorityVoting(args.parts, name)
  
      with jt.no_grad():
          for meshes, labels, mesh_infos in tqdm(dataset, desc=str(epoch)):
              mesh_tensor = to_mesh_tensor(meshes)
              mesh_labels = jt.int32(labels)
              outputs = net(mesh_tensor)
              preds = np.argmax(outputs.data, axis=1)
  
              batch_acc = (labels == preds).sum(axis=1) / meshes['Fs']
              batch_oacc = compute_original_accuracy(mesh_infos, preds, mesh_labels)
              acc += np.sum(batch_acc)
              oacc += np.sum(batch_oacc)
              update_label_accuracy(preds, mesh_labels, label_acc)
              voted.vote(mesh_infos, preds, mesh_labels)
  
      acc /= dataset.total_len
      oacc /= dataset.total_len
      voacc = voted.compute_accuracy(save_results=True)
      print('test acc = ', acc)
      print('test acc [original] = ', oacc)
      print('test acc [original] [voted] = ', voacc)
      print('test acc per label = ', label_acc / dataset.total_len)
      writer.add_scalar('test-acc', acc, global_step=epoch)
      writer.add_scalar('test-oacc', oacc, global_step=epoch)
      writer.add_scalar('test-voacc', voacc, global_step=epoch)
  
      if test.best_oacc < oacc:
          test.best_oacc = oacc
          net.save(os.path.join('checkpoints', name, f'oacc-{oacc:.4f}.pkl'))
  
      if test.best_voacc < voacc:
          test.best_voacc = voacc
          net.save(os.path.join('checkpoints', name, f'voacc-{voacc:.4f}.pkl'))
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('mode', choices=['train', 'test'])
      parser.add_argument('--name', type=str, required=True)
      parser.add_argument('--dataroot', type=str, required=True)
      parser.add_argument('--batch_size', type=int, default=8)
      parser.add_argument('--optim', choices=['adam', 'sgd'], default='adam')
      parser.add_argument('--lr', type=float, default=2e-2)
      parser.add_argument('--weight_decay', type=float, default=0)
      parser.add_argument('--checkpoint', type=str)
      parser.add_argument('--upsample', choices=['nearest', 'bilinear'], default='bilinear')
      parser.add_argument('--parts', type=int, default=8)
      parser.add_argument('--augment_scale', action='store_true')
      parser.add_argument('--augment_orient', action='store_true')
      parser.add_argument('--arch', choices=['unet', 'deeplab', 'vunet'], default='unet')
      parser.add_argument('--backbone', choices=['resnet18', 'resnet50'], default='resnet50')
      parser.add_argument('--globalpool', choices=['max', 'mean'], default='mean')
  
      args = parser.parse_args()
      mode = args.mode
      name = args.name
      batch_size = args.batch_size
      dataroot = args.dataroot
  
      net = None
      if args.arch == 'deeplab':
          net = MeshDeepLab(13, args.parts, args.backbone, globalpool=args.globalpool)
      elif args.arch == 'unet':
          net = MeshVanillaUnet(13, args.parts, upsample=args.upsample)
      
      if args.optim == 'adam':
          optim = Adam(net.parameters(), lr=args.lr, weight_decay=args.weight_decay)
      else:
          optim = SGD(net.parameters(), lr=args.lr, weight_decay=args.weight_decay)
      scheduler = MultiStepLR(optim, milestones=[50, 100, 150], gamma=0.1)
   
      writer = SummaryWriter("logs/" + name)
      print('name:', name)
  
      augments = []
      if args.augment_scale:
          augments.append('scale')
      if args.augment_orient:
          augments.append('orient')
  
      if mode == 'train':
          train_dataset = SegmentationDataset(dataroot, batch_size=batch_size, 
              shuffle=True, train=True, num_workers=4, augments=augments)
      test_dataset = SegmentationDataset(dataroot, batch_size=8, shuffle=False, 
          train=False, num_workers=4)
  
      checkpoint_path = os.path.join('checkpoints', name)
      checkpoint_name = os.path.join(checkpoint_path, name + '-latest.pkl')
      os.makedirs(checkpoint_path, exist_ok=True)
  
      if args.checkpoint is not None:
          net.load(args.checkpoint)
  
      train.step = 0
      test.best_oacc = 0
      test.best_voacc = 0
  
      if args.mode == 'train':
          for epoch in range(500):
              train(net, optim, train_dataset, writer, epoch)
              test(net, test_dataset, writer, epoch, args)
              scheduler.step()
              net.save(checkpoint_name)
      else:
          test(net, test_dataset, writer, 0, args)
  