code: def train(args, model, train_dataloader, idx, class_name):
      train_outputs = OrderedDict(
          [("layer1", []), ("layer2", []), ("layer3", [])]
      )
  
      # extract train set features
      epoch_begin = time.time()
      end_time = time.time()
  
      for index, x in enumerate(train_dataloader):
          start_time = time.time()
          data_time = start_time - end_time
  
          # model prediction
          with paddle.no_grad():
              outputs = model(x)
  
          # get intermediate layer outputs
          for k, v in zip(train_outputs.keys(), outputs):
              train_outputs[k].append(v.cpu().detach())
  
          end_time = time.time()
          batch_time = end_time - start_time
  
          if index % args.print_freq == 0:
              print(
                  datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                  + "\t"
                  + "Epoch {}[{}/{}]: loss:{:.5f}, lr:{:.5f}, batch time:{:.4f}, data time:{:.4f}".format(
                      0,
                      index + 1,
                      len(train_dataloader),
                      0,
                      float(0),
                      float(batch_time),
                      float(data_time),
                  )
              )
  
      for k, v in train_outputs.items():
          train_outputs[k] = paddle.concat(v, 0)
  
      # Embedding concat
      embedding_vectors = train_outputs["layer1"]
      for layer_name in ["layer2", "layer3"]:
          layer_embedding = train_outputs[layer_name]
          layer_embedding = F.interpolate(
              layer_embedding,
              size=embedding_vectors.shape[-2:],
              mode="nearest",
          )
          embedding_vectors = paddle.concat(
              (embedding_vectors, layer_embedding), 1
          )
  
      # randomly select d dimension
      embedding_vectors = paddle.index_select(embedding_vectors, idx, 1)
      # calculate multivariate Gaussian distribution
      B, C, H, W = embedding_vectors.shape
      embedding_vectors = embedding_vectors.reshape((B, C, H * W))
      mean = paddle.mean(embedding_vectors, axis=0).numpy()
      cov = paddle.zeros((C, C, H * W)).numpy()
      I = np.identity(C)
      for i in range(H * W):
          cov[:, :, i] = (
              np.cov(embedding_vectors[:, :, i].numpy(), rowvar=False)
              + 0.01 * I
          )
      # save learned distribution
      train_outputs = [mean, cov]
      model.distribution = train_outputs
      t = time.time() - epoch_begin
      print(
          datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
          + "\t"
          + "Train ends, total {:.2f}s".format(0, t)
      )
      if args.save_model:
          print(
              datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
              + "\t"
              + "Saving model..."
          )
          save_name = os.path.join(
              args.save_path,
              args.backbone,
              args.category,
              "{}.pdparams".format(class_name),
          )
          dir_name = os.path.dirname(save_name)
          if dir_name and not os.path.exists(dir_name):
              os.makedirs(dir_name)
          state_dict = {
              "params": model.model.state_dict(),
              "distribution": model.distribution,
          }
          paddle.save(state_dict, save_name)
          print(
              datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
              + "\t"
              + "Save model in {}".format(str(save_name))
          )
  
file_path: PaddleSeg/contrib/QualityInspector/tools/uad/padim/train.py
imports: import os
  import sys
  import time
  import random
  import argparse
  import datetime
  from random import sample
  from collections import OrderedDict
  import numpy
  import pandas
  import paddle
  import paddle.nn.functional
  from paddle.io import DataLoader
  import qinspector.uad.datasets.mvtec
  from qinspector.uad.models.padim import ResNet_PaDiM
  from qinspector.cvlib.uad_configs import ConfigParser
  from val import val
label_desc: Use image node_id 3 for calling a global function with example usage: train(args, model, train_dataloader, idx, class_name) without return types