cls1 basic

基本写法

  1. import numpy as np
  3. import torch
  4. from torch import nn
  6. from detectron2.data.build import TrainingSampler
  8. from pyxlpr.d2 import D2Trainer
  11. class ValueCls(D2Trainer):
  12.     """ 输入数据是3个[1, 10]的整数，做分类任务，判断和是否在(10,20)之间 """
  14.     @classmethod
  15.     def build_model(cls, cfg):
  16.         """ 自定义模型 """
  17.         device = torch.device(cfg.MODEL.DEVICE)
  19.         class ParityModel(nn.Module):
  20.             def __init__(self):
  21.                 super().__init__()
  23.                 self.classifier = nn.Sequential(
  24.                     nn.Linear(in_features=3, out_features=2),
  25.                     nn.Sigmoid(),
  26.                     nn.Linear(in_features=2, out_features=2),
  27.                     nn.Sigmoid(),
  28.                     nn.Linear(in_features=2, out_features=1),
  29.                     nn.Sigmoid(),
  30.                 )
  31.                 self.criteon = nn.BCELoss()  # 二分类交叉熵损失
  33.             def forward(self, batched_inputs):
  34.                 x = batched_inputs[0].type(torch.FloatTensor).to(device)
  35.                 logits = self.classifier(x)
  37.                 if self.training:
  38.                     y = batched_inputs[1].unsqueeze(-1).type(torch.FloatTensor).to(device)
  39.                     loss = self.criteon(logits, y)
  40.                     return {'loss': loss}  # 损失要以字典的形式返回
  41.                 else:
  42.                     y_hat = (logits > 0.5).type(torch.int)
  43.                     return y_hat
  45.         return ParityModel().to(device)
  47.     @classmethod
  48.     def build_train_loader(cls, cfg):
  49.         """ 自定义训练集 """
  50.         n = 5000  # 数据量
  51.         data = np.random.randint(1, 10, [n, 3])  # 数据
  52.         dataloader = torch.utils.data.DataLoader([(x, 10 < sum(x) < 20) for x in data],
  53.                                                  sampler=TrainingSampler(n),  # 无限取用的数据流，避免一个epoch完迭代结束
  54.                                                  batch_size=cfg.SOLVER.IMS_PER_BATCH)
  55.         return dataloader
  57.     @classmethod
  58.     def build_test_loader(cls, cfg, dataset_name):
  59.         """ 自定义验证集 """
  60.         n = 100
  61.         data = np.random.randint(1, 10, [n, 3])
  62.         dataloader = torch.utils.data.DataLoader([(x, 10 < sum(x) < 20) for x in data],
  63.                                                  batch_size=cfg.SOLVER.IMS_PER_BATCH)
  64.         return dataloader
  66.     @classmethod
  67.     def build_evaluator(cls, cfg, dataset_name, output_folder=None):
  68.         """ 自定义测评方法 """
  69.         from detectron2.evaluation import DatasetEvaluator
  71.         class Eval(DatasetEvaluator):
  72.             def reset(self):
  73.                 self.acc = 0  # 分类正确的样本数
  74.                 self.total = 0  # 总样本数
  76.             def process(self, inputs, outputs):
  77.                 """
  78.                 :param inputs: dataloader每次batch得到的内容结构
  79.                 :param outputs: model.inference对应的返回值
  80.                 :return:
  81.                 """
  82.                 self.total += len(outputs)
  83.                 self.acc += int(sum(inputs[1].unsqueeze(-1) == outputs)[0])
  85.             def evaluate(self):
  86.                 # 结果要以字典的形式返回，第一层是任务名，第二层是多个指标结果
  87.                 return {'Classification': {'Accuracy': self.acc / self.total, 'total': self.total}}
  89.         return Eval()
  92. CONFIGS = r"""
  93. MODEL:
  94.   DEVICE: 'cpu'
  95. SOLVER:
  96.   IMS_PER_BATCH: 16
  97.   MAX_ITER: 3125
  98.   BASE_LR: 0.01
  99. OUTPUT_DIR: "/home/chenkunze/det2/valuecls"
  100. """
  102. if __name__ == "__main__":
  103.     np.random.seed(4101)
  104.     torch.manual_seed(4102)
  106.     ValueCls(CONFIGS).launch()
  107.     # 目前精度：0.6100
  108.     # 框架不一样，相比于原生pytorch，跑出来确实有些不同吧

cls2 enhance

加强版