classifier.py 33 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798
  1. # Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. import math
  16. import os.path as osp
  17. from collections import OrderedDict
  18. import numpy as np
  19. import paddle
  20. from paddle import to_tensor
  21. import paddle.nn.functional as F
  22. from paddle.static import InputSpec
  23. from paddlex.utils import logging, TrainingStats, DisablePrint
  24. from paddlex.cv.models.base import BaseModel
  25. from paddlex.cv.transforms import arrange_transforms
  26. from paddlex.cv.transforms.operators import Resize
  27. with DisablePrint():
  28. from PaddleClas.ppcls.modeling import architectures
  29. from PaddleClas.ppcls.modeling.loss import CELoss
  30. __all__ = [
  31. "ResNet18", "ResNet34", "ResNet50", "ResNet101", "ResNet152",
  32. "ResNet18_vd", "ResNet34_vd", "ResNet50_vd", "ResNet50_vd_ssld",
  33. "ResNet101_vd", "ResNet101_vd_ssld", "ResNet152_vd", "ResNet200_vd",
  34. "AlexNet", "DarkNet53", "MobileNetV1", "MobileNetV2", "MobileNetV3_small",
  35. "MobileNetV3_small_ssld", "MobileNetV3_large", "MobileNetV3_large_ssld",
  36. "DenseNet121", "DenseNet161", "DenseNet169", "DenseNet201", "DenseNet264",
  37. "HRNet_W18_C", "HRNet_W30_C", "HRNet_W32_C", "HRNet_W40_C", "HRNet_W44_C",
  38. "HRNet_W48_C", "HRNet_W64_C", "Xception41", "Xception65", "Xception71",
  39. "ShuffleNetV2", "ShuffleNetV2_swish"
  40. ]
  41. class BaseClassifier(BaseModel):
  42. """Parent class of all classification models.
  43. Args:
  44. model_name (str, optional): Name of classification model. Defaults to 'ResNet50'.
  45. num_classes (int, optional): The number of target classes. Defaults to 1000.
  46. """
  47. def __init__(self, model_name='ResNet50', num_classes=1000, **params):
  48. self.init_params = locals()
  49. self.init_params.update(params)
  50. if 'lr_mult_list' in self.init_params:
  51. del self.init_params['lr_mult_list']
  52. super(BaseClassifier, self).__init__('classifier')
  53. if not hasattr(architectures, model_name):
  54. raise Exception("ERROR: There's no model named {}.".format(
  55. model_name))
  56. self.model_name = model_name
  57. self.labels = None
  58. self.num_classes = num_classes
  59. for k, v in params.items():
  60. setattr(self, k, v)
  61. self.net = self.build_net(**params)
  62. def build_net(self, **params):
  63. with paddle.utils.unique_name.guard():
  64. net = architectures.__dict__[self.model_name](
  65. class_dim=self.num_classes, **params)
  66. return net
  67. def _fix_transforms_shape(self, image_shape):
  68. if hasattr(self, 'test_transforms'):
  69. if self.test_transforms is not None:
  70. self.test_transforms.transforms.append(
  71. Resize(target_size=image_shape))
  72. def _get_test_inputs(self, image_shape):
  73. if image_shape is not None:
  74. if len(image_shape) == 2:
  75. image_shape = [1, 3] + image_shape
  76. self._fix_transforms_shape(image_shape[-2:])
  77. else:
  78. image_shape = [None, 3, -1, -1]
  79. self.fixed_input_shape = image_shape
  80. input_spec = [
  81. InputSpec(
  82. shape=image_shape, name='image', dtype='float32')
  83. ]
  84. return input_spec
  85. def run(self, net, inputs, mode):
  86. net_out = net(inputs[0])
  87. softmax_out = F.softmax(net_out)
  88. if mode == 'test':
  89. outputs = OrderedDict([('prediction', softmax_out)])
  90. elif mode == 'eval':
  91. pred = softmax_out
  92. gt = inputs[1]
  93. labels = inputs[1].reshape([-1, 1])
  94. acc1 = paddle.metric.accuracy(softmax_out, label=labels)
  95. k = min(5, self.num_classes)
  96. acck = paddle.metric.accuracy(softmax_out, label=labels, k=k)
  97. # multi cards eval
  98. if paddle.distributed.get_world_size() > 1:
  99. acc1 = paddle.distributed.all_reduce(
  100. acc1, op=paddle.distributed.ReduceOp.
  101. SUM) / paddle.distributed.get_world_size()
  102. acck = paddle.distributed.all_reduce(
  103. acck, op=paddle.distributed.ReduceOp.
  104. SUM) / paddle.distributed.get_world_size()
  105. pred = list()
  106. gt = list()
  107. paddle.distributed.all_gather(pred, softmax_out)
  108. paddle.distributed.all_gather(gt, inputs[1])
  109. pred = paddle.concat(pred, axis=0)
  110. gt = paddle.concat(gt, axis=0)
  111. outputs = OrderedDict([('acc1', acc1), ('acc{}'.format(k), acck),
  112. ('prediction', pred), ('labels', gt)])
  113. else:
  114. # mode == 'train'
  115. labels = inputs[1].reshape([-1, 1])
  116. loss = CELoss(class_dim=self.num_classes)
  117. loss = loss(net_out, inputs[1])
  118. acc1 = paddle.metric.accuracy(softmax_out, label=labels, k=1)
  119. k = min(5, self.num_classes)
  120. acck = paddle.metric.accuracy(softmax_out, label=labels, k=k)
  121. outputs = OrderedDict([('loss', loss), ('acc1', acc1),
  122. ('acc{}'.format(k), acck)])
  123. return outputs
  124. def default_optimizer(self, parameters, learning_rate, warmup_steps,
  125. warmup_start_lr, lr_decay_epochs, lr_decay_gamma,
  126. num_steps_each_epoch):
  127. boundaries = [b * num_steps_each_epoch for b in lr_decay_epochs]
  128. values = [
  129. learning_rate * (lr_decay_gamma**i)
  130. for i in range(len(lr_decay_epochs) + 1)
  131. ]
  132. scheduler = paddle.optimizer.lr.PiecewiseDecay(boundaries, values)
  133. if warmup_steps > 0:
  134. if warmup_steps > lr_decay_epochs[0] * num_steps_each_epoch:
  135. logging.error(
  136. "In function train(), parameters should satisfy: "
  137. "warmup_steps <= lr_decay_epochs[0]*num_samples_in_train_dataset",
  138. exit=False)
  139. logging.error(
  140. "See this doc for more information: "
  141. "https://github.com/PaddlePaddle/PaddleX/blob/develop/docs/appendix/parameters.md#notice",
  142. exit=False)
  143. logging.error(
  144. "warmup_steps should less than {} or lr_decay_epochs[0] greater than {}, "
  145. "please modify 'lr_decay_epochs' or 'warmup_steps' in train function".
  146. format(lr_decay_epochs[0] * num_steps_each_epoch,
  147. warmup_steps // num_steps_each_epoch))
  148. scheduler = paddle.optimizer.lr.LinearWarmup(
  149. learning_rate=scheduler,
  150. warmup_steps=warmup_steps,
  151. start_lr=warmup_start_lr,
  152. end_lr=learning_rate)
  153. optimizer = paddle.optimizer.Momentum(
  154. scheduler,
  155. momentum=.9,
  156. weight_decay=paddle.regularizer.L2Decay(coeff=1e-04),
  157. parameters=parameters)
  158. return optimizer
  159. def train(self,
  160. num_epochs,
  161. train_dataset,
  162. train_batch_size=64,
  163. eval_dataset=None,
  164. optimizer=None,
  165. save_interval_epochs=1,
  166. log_interval_steps=10,
  167. save_dir='output',
  168. pretrain_weights='IMAGENET',
  169. learning_rate=.025,
  170. warmup_steps=0,
  171. warmup_start_lr=0.0,
  172. lr_decay_epochs=(30, 60, 90),
  173. lr_decay_gamma=0.1,
  174. early_stop=False,
  175. early_stop_patience=5,
  176. use_vdl=True):
  177. """
  178. Train the model.
  179. Args:
  180. num_epochs(int): The number of epochs.
  181. train_dataset(paddlex.dataset): Training dataset.
  182. train_batch_size(int, optional): Total batch size among all cards used in training. Defaults to 64.
  183. eval_dataset(paddlex.dataset, optional):
  184. Evaluation dataset. If None, the model will not be evaluated during training process. Defaults to None.
  185. optimizer(paddle.optimizer.Optimizer or None, optional):
  186. Optimizer used for training. If None, a default optimizer is used. Defaults to None.
  187. save_interval_epochs(int, optional): Epoch interval for saving the model. Defaults to 1.
  188. log_interval_steps(int, optional): Step interval for printing training information. Defaults to 10.
  189. save_dir(str, optional): Directory to save the model. Defaults to 'output'.
  190. pretrain_weights(str or None, optional):
  191. None or name/path of pretrained weights. If None, no pretrained weights will be loaded. Defaults to 'IMAGENET'.
  192. learning_rate(float, optional): Learning rate for training. Defaults to .025.
  193. warmup_steps(int, optional): The number of steps of warm-up training. Defaults to 0.
  194. warmup_start_lr(float, optional): Start learning rate of warm-up training. Defaults to 0..
  195. lr_decay_epochs(List[int] or Tuple[int], optional):
  196. Epoch milestones for learning rate decay. Defaults to (20, 60, 90).
  197. lr_decay_gamma(float, optional): Gamma coefficient of learning rate decay, default .1.
  198. early_stop(bool, optional): Whether to adopt early stop strategy. Defaults to False.
  199. early_stop_patience(int, optional): Early stop patience. Defaults to 5.
  200. use_vdl(bool, optional): Whether to use VisualDL to monitor the training process. Defaults to True.
  201. """
  202. self.labels = train_dataset.labels
  203. # build optimizer if not defined
  204. if optimizer is None:
  205. num_steps_each_epoch = len(train_dataset) // train_batch_size
  206. self.optimizer = self.default_optimizer(
  207. parameters=self.net.parameters(),
  208. learning_rate=learning_rate,
  209. warmup_steps=warmup_steps,
  210. warmup_start_lr=warmup_start_lr,
  211. lr_decay_epochs=lr_decay_epochs,
  212. lr_decay_gamma=lr_decay_gamma,
  213. num_steps_each_epoch=num_steps_each_epoch)
  214. else:
  215. self.optimizer = optimizer
  216. # initiate weights
  217. if pretrain_weights is not None and not osp.exists(pretrain_weights):
  218. if pretrain_weights not in ['IMAGENET']:
  219. logging.warning(
  220. "Path of pretrain_weights('{}') does not exist!".format(
  221. pretrain_weights))
  222. logging.warning(
  223. "Pretrain_weights is forcibly set to 'IMAGENET'. "
  224. "If don't want to use pretrain weights, "
  225. "set pretrain_weights to be None.")
  226. pretrain_weights = 'IMAGENET'
  227. elif pretrain_weights is not None and osp.exists(pretrain_weights):
  228. if osp.splitext(pretrain_weights)[-1] != '.pdparams':
  229. logging.error(
  230. "Invalid pretrain weights. Please specify a '.pdparams' file.",
  231. exit=True)
  232. pretrained_dir = osp.join(save_dir, 'pretrain')
  233. self.net_initialize(
  234. pretrain_weights=pretrain_weights, save_dir=pretrained_dir)
  235. # start train loop
  236. self.train_loop(
  237. num_epochs=num_epochs,
  238. train_dataset=train_dataset,
  239. train_batch_size=train_batch_size,
  240. eval_dataset=eval_dataset,
  241. save_interval_epochs=save_interval_epochs,
  242. log_interval_steps=log_interval_steps,
  243. save_dir=save_dir,
  244. early_stop=early_stop,
  245. early_stop_patience=early_stop_patience,
  246. use_vdl=use_vdl)
  247. def quant_aware_train(self,
  248. num_epochs,
  249. train_dataset,
  250. train_batch_size=64,
  251. eval_dataset=None,
  252. optimizer=None,
  253. save_interval_epochs=1,
  254. log_interval_steps=10,
  255. save_dir='output',
  256. learning_rate=.000025,
  257. warmup_steps=0,
  258. warmup_start_lr=0.0,
  259. lr_decay_epochs=(30, 60, 90),
  260. lr_decay_gamma=0.1,
  261. early_stop=False,
  262. early_stop_patience=5,
  263. use_vdl=True,
  264. quant_config=None):
  265. """
  266. Quantization-aware training.
  267. Args:
  268. num_epochs(int): The number of epochs.
  269. train_dataset(paddlex.dataset): Training dataset.
  270. train_batch_size(int, optional): Total batch size among all cards used in training. Defaults to 64.
  271. eval_dataset(paddlex.dataset, optional):
  272. Evaluation dataset. If None, the model will not be evaluated during training process. Defaults to None.
  273. optimizer(paddle.optimizer.Optimizer or None, optional):
  274. Optimizer used for training. If None, a default optimizer is used. Defaults to None.
  275. save_interval_epochs(int, optional): Epoch interval for saving the model. Defaults to 1.
  276. log_interval_steps(int, optional): Step interval for printing training information. Defaults to 10.
  277. save_dir(str, optional): Directory to save the model. Defaults to 'output'.
  278. learning_rate(float, optional): Learning rate for training. Defaults to .025.
  279. warmup_steps(int, optional): The number of steps of warm-up training. Defaults to 0.
  280. warmup_start_lr(float, optional): Start learning rate of warm-up training. Defaults to 0..
  281. lr_decay_epochs(List[int] or Tuple[int], optional):
  282. Epoch milestones for learning rate decay. Defaults to (20, 60, 90).
  283. lr_decay_gamma(float, optional): Gamma coefficient of learning rate decay, default .1.
  284. early_stop(bool, optional): Whether to adopt early stop strategy. Defaults to False.
  285. early_stop_patience(int, optional): Early stop patience. Defaults to 5.
  286. use_vdl(bool, optional): Whether to use VisualDL to monitor the training process. Defaults to True.
  287. quant_config(dict or None, optional): Quantization configuration. If None, a default rule of thumb
  288. configuration will be used. Defaults to None.
  289. """
  290. self._prepare_qat(quant_config)
  291. self.train(
  292. num_epochs=num_epochs,
  293. train_dataset=train_dataset,
  294. train_batch_size=train_batch_size,
  295. eval_dataset=eval_dataset,
  296. optimizer=optimizer,
  297. save_interval_epochs=save_interval_epochs,
  298. log_interval_steps=log_interval_steps,
  299. save_dir=save_dir,
  300. pretrain_weights=None,
  301. learning_rate=learning_rate,
  302. warmup_steps=warmup_steps,
  303. warmup_start_lr=warmup_start_lr,
  304. lr_decay_epochs=lr_decay_epochs,
  305. lr_decay_gamma=lr_decay_gamma,
  306. early_stop=early_stop,
  307. early_stop_patience=early_stop_patience,
  308. use_vdl=use_vdl)
  309. def evaluate(self, eval_dataset, batch_size=1, return_details=False):
  310. """
  311. Evaluate the model.
  312. Args:
  313. eval_dataset(paddlex.dataset): Evaluation dataset.
  314. batch_size(int, optional): Total batch size among all cards used for evaluation. Defaults to 1.
  315. return_details(bool, optional): Whether to return evaluation details. Defaults to False.
  316. Returns:
  317. collections.OrderedDict with key-value pairs: {"acc1": `top 1 accuracy`, "acc5": `top 5 accuracy`}.
  318. """
  319. # 给transform添加arrange操作
  320. arrange_transforms(
  321. model_type=self.model_type,
  322. transforms=eval_dataset.transforms,
  323. mode='eval')
  324. self.net.eval()
  325. nranks = paddle.distributed.get_world_size()
  326. local_rank = paddle.distributed.get_rank()
  327. if nranks > 1:
  328. # Initialize parallel environment if not done.
  329. if not paddle.distributed.parallel.parallel_helper._is_parallel_ctx_initialized(
  330. ):
  331. paddle.distributed.init_parallel_env()
  332. self.eval_data_loader = self.build_data_loader(
  333. eval_dataset, batch_size=batch_size, mode='eval')
  334. eval_metrics = TrainingStats()
  335. if return_details:
  336. true_labels = list()
  337. pred_scores = list()
  338. logging.info(
  339. "Start to evaluate(total_samples={}, total_steps={})...".format(
  340. eval_dataset.num_samples,
  341. math.ceil(eval_dataset.num_samples * 1.0 / batch_size)))
  342. with paddle.no_grad():
  343. for step, data in enumerate(self.eval_data_loader()):
  344. outputs = self.run(self.net, data, mode='eval')
  345. if return_details:
  346. true_labels.extend(outputs['labels'].tolist())
  347. pred_scores.extend(outputs['prediction'].tolist())
  348. outputs.pop('prediction')
  349. outputs.pop('labels')
  350. eval_metrics.update(outputs)
  351. if return_details:
  352. eval_details = {
  353. 'true_labels': true_labels,
  354. 'pred_scores': pred_scores
  355. }
  356. return eval_metrics.get(), eval_details
  357. else:
  358. return eval_metrics.get()
  359. def predict(self, img_file, transforms=None, topk=1):
  360. """
  361. Do inference.
  362. Args:
  363. img_file(List[np.ndarray or str], str or np.ndarray): img_file(list or str or np.array):
  364. Image path or decoded image data in a BGR format, which also could constitute a list,
  365. meaning all images to be predicted as a mini-batch.
  366. transforms(paddlex.transforms.Compose or None, optional):
  367. Transforms for inputs. If None, the transforms for evaluation process will be used. Defaults to None.
  368. topk(int, optional): Keep topk results in prediction. Defaults to 1.
  369. Returns:
  370. If img_file is a string or np.array, the result is a dict with key-value pairs:
  371. {"category_id": `category_id`, "category": `category`, "score": `score`}.
  372. If img_file is a list, the result is a list composed of dicts with the corresponding fields:
  373. category_id(int): the predicted category ID
  374. category(str): category name
  375. score(float): confidence
  376. """
  377. if transforms is None and not hasattr(self, 'test_transforms'):
  378. raise Exception("transforms need to be defined, now is None.")
  379. if transforms is None:
  380. transforms = self.test_transforms
  381. true_topk = min(self.num_classes, topk)
  382. if isinstance(img_file, (str, np.ndarray)):
  383. images = [img_file]
  384. else:
  385. images = img_file
  386. im = self._preprocess(images, transforms, self.model_type)
  387. self.net.eval()
  388. with paddle.no_grad():
  389. outputs = self.run(self.net, im, mode='test')
  390. prediction = outputs['prediction'].numpy()
  391. prediction = self._postprocess(prediction, true_topk, self.labels)
  392. if isinstance(img_file, (str, np.ndarray)):
  393. prediction = prediction[0]
  394. return prediction
  395. def _preprocess(self, images, transforms, model_type):
  396. arrange_transforms(
  397. model_type=model_type, transforms=transforms, mode='test')
  398. batch_im = list()
  399. for im in images:
  400. sample = {'image': im}
  401. batch_im.append(transforms(sample))
  402. batch_im = to_tensor(batch_im)
  403. return batch_im,
  404. def _postprocess(self, results, true_topk, labels):
  405. preds = list()
  406. for i, pred in enumerate(results):
  407. pred_label = np.argsort(pred)[::-1][:true_topk]
  408. preds.append([{
  409. 'category_id': l,
  410. 'category': labels[l],
  411. 'score': results[i][l]
  412. } for l in pred_label])
  413. return preds
  414. class ResNet18(BaseClassifier):
  415. def __init__(self, num_classes=1000):
  416. super(ResNet18, self).__init__(
  417. model_name='ResNet18', num_classes=num_classes)
  418. class ResNet34(BaseClassifier):
  419. def __init__(self, num_classes=1000):
  420. super(ResNet34, self).__init__(
  421. model_name='ResNet34', num_classes=num_classes)
  422. class ResNet50(BaseClassifier):
  423. def __init__(self, num_classes=1000):
  424. super(ResNet50, self).__init__(
  425. model_name='ResNet50', num_classes=num_classes)
  426. class ResNet101(BaseClassifier):
  427. def __init__(self, num_classes=1000):
  428. super(ResNet101, self).__init__(
  429. model_name='ResNet101', num_classes=num_classes)
  430. class ResNet152(BaseClassifier):
  431. def __init__(self, num_classes=1000):
  432. super(ResNet152, self).__init__(
  433. model_name='ResNet152', num_classes=num_classes)
  434. class ResNet18_vd(BaseClassifier):
  435. def __init__(self, num_classes=1000):
  436. super(ResNet18_vd, self).__init__(
  437. model_name='ResNet18_vd', num_classes=num_classes)
  438. class ResNet34_vd(BaseClassifier):
  439. def __init__(self, num_classes=1000):
  440. super(ResNet34_vd, self).__init__(
  441. model_name='ResNet34_vd', num_classes=num_classes)
  442. class ResNet50_vd(BaseClassifier):
  443. def __init__(self, num_classes=1000):
  444. super(ResNet50_vd, self).__init__(
  445. model_name='ResNet50_vd', num_classes=num_classes)
  446. class ResNet50_vd_ssld(BaseClassifier):
  447. def __init__(self, num_classes=1000):
  448. super(ResNet50_vd_ssld, self).__init__(
  449. model_name='ResNet50_vd',
  450. num_classes=num_classes,
  451. lr_mult_list=[.1, .1, .2, .2, .3])
  452. self.model_name = 'ResNet50_vd_ssld'
  453. class ResNet101_vd(BaseClassifier):
  454. def __init__(self, num_classes=1000):
  455. super(ResNet101_vd, self).__init__(
  456. model_name='ResNet101_vd', num_classes=num_classes)
  457. class ResNet101_vd_ssld(BaseClassifier):
  458. def __init__(self, num_classes=1000):
  459. super(ResNet101_vd_ssld, self).__init__(
  460. model_name='ResNet101_vd',
  461. num_classes=num_classes,
  462. lr_mult_list=[.1, .1, .2, .2, .3])
  463. self.model_name = 'ResNet101_vd_ssld'
  464. class ResNet152_vd(BaseClassifier):
  465. def __init__(self, num_classes=1000):
  466. super(ResNet152_vd, self).__init__(
  467. model_name='ResNet152_vd', num_classes=num_classes)
  468. class ResNet200_vd(BaseClassifier):
  469. def __init__(self, num_classes=1000):
  470. super(ResNet200_vd, self).__init__(
  471. model_name='ResNet200_vd', num_classes=num_classes)
  472. class AlexNet(BaseClassifier):
  473. def __init__(self, num_classes=1000):
  474. super(AlexNet, self).__init__(
  475. model_name='AlexNet', num_classes=num_classes)
  476. def _get_test_inputs(self, image_shape):
  477. if image_shape is not None:
  478. if len(image_shape) == 2:
  479. image_shape = [None, 3] + image_shape
  480. else:
  481. image_shape = [None, 3, 224, 224]
  482. logging.warning(
  483. '[Important!!!] When exporting inference model for {},'.format(
  484. self.__class__.__name__) +
  485. ' if fixed_input_shape is not set, it will be forcibly set to [None, 3, 224, 224]'
  486. +
  487. 'Please check image shape after transforms is [3, 224, 224], if not, fixed_input_shape '
  488. + 'should be specified manually.')
  489. self._fix_transforms_shape(image_shape[-2:])
  490. self.fixed_input_shape = image_shape
  491. input_spec = [
  492. InputSpec(
  493. shape=image_shape, name='image', dtype='float32')
  494. ]
  495. return input_spec
  496. class DarkNet53(BaseClassifier):
  497. def __init__(self, num_classes=1000):
  498. super(DarkNet53, self).__init__(
  499. model_name='DarkNet53', num_classes=num_classes)
  500. class MobileNetV1(BaseClassifier):
  501. def __init__(self, num_classes=1000, scale=1.0):
  502. supported_scale = [.25, .5, .75, 1.0]
  503. if scale not in supported_scale:
  504. logging.warning("scale={} is not supported by MobileNetV1, "
  505. "scale is forcibly set to 1.0".format(scale))
  506. scale = 1.0
  507. if scale == 1:
  508. model_name = 'MobileNetV1'
  509. else:
  510. model_name = 'MobileNetV1_x' + str(scale).replace('.', '_')
  511. self.scale = scale
  512. super(MobileNetV1, self).__init__(
  513. model_name=model_name, num_classes=num_classes)
  514. class MobileNetV2(BaseClassifier):
  515. def __init__(self, num_classes=1000, scale=1.0):
  516. supported_scale = [.25, .5, .75, 1.0, 1.5, 2.0]
  517. if scale not in supported_scale:
  518. logging.warning("scale={} is not supported by MobileNetV2, "
  519. "scale is forcibly set to 1.0".format(scale))
  520. scale = 1.0
  521. if scale == 1:
  522. model_name = 'MobileNetV2'
  523. else:
  524. model_name = 'MobileNetV2_x' + str(scale).replace('.', '_')
  525. super(MobileNetV2, self).__init__(
  526. model_name=model_name, num_classes=num_classes)
  527. class MobileNetV3_small(BaseClassifier):
  528. def __init__(self, num_classes=1000, scale=1.0):
  529. supported_scale = [.35, .5, .75, 1.0, 1.25]
  530. if scale not in supported_scale:
  531. logging.warning("scale={} is not supported by MobileNetV3_small, "
  532. "scale is forcibly set to 1.0".format(scale))
  533. scale = 1.0
  534. model_name = 'MobileNetV3_small_x' + str(float(scale)).replace('.',
  535. '_')
  536. super(MobileNetV3_small, self).__init__(
  537. model_name=model_name, num_classes=num_classes)
  538. class MobileNetV3_small_ssld(BaseClassifier):
  539. def __init__(self, num_classes=1000, scale=1.0):
  540. supported_scale = [.35, 1.0]
  541. if scale not in supported_scale:
  542. logging.warning(
  543. "scale={} is not supported by MobileNetV3_small_ssld, "
  544. "scale is forcibly set to 1.0".format(scale))
  545. scale = 1.0
  546. model_name = 'MobileNetV3_small_x' + str(float(scale)).replace('.',
  547. '_')
  548. super(MobileNetV3_small_ssld, self).__init__(
  549. model_name=model_name, num_classes=num_classes)
  550. self.model_name = model_name + '_ssld'
  551. class MobileNetV3_large(BaseClassifier):
  552. def __init__(self, num_classes=1000, scale=1.0):
  553. supported_scale = [.35, .5, .75, 1.0, 1.25]
  554. if scale not in supported_scale:
  555. logging.warning("scale={} is not supported by MobileNetV3_large, "
  556. "scale is forcibly set to 1.0".format(scale))
  557. scale = 1.0
  558. model_name = 'MobileNetV3_large_x' + str(float(scale)).replace('.',
  559. '_')
  560. super(MobileNetV3_large, self).__init__(
  561. model_name=model_name, num_classes=num_classes)
  562. class MobileNetV3_large_ssld(BaseClassifier):
  563. def __init__(self, num_classes=1000):
  564. super(MobileNetV3_large_ssld, self).__init__(
  565. model_name='MobileNetV3_large_x1_0', num_classes=num_classes)
  566. self.model_name = 'MobileNetV3_large_x1_0_ssld'
  567. class DenseNet121(BaseClassifier):
  568. def __init__(self, num_classes=1000):
  569. super(DenseNet121, self).__init__(
  570. model_name='DenseNet121', num_classes=num_classes)
  571. class DenseNet161(BaseClassifier):
  572. def __init__(self, num_classes=1000):
  573. super(DenseNet161, self).__init__(
  574. model_name='DenseNet161', num_classes=num_classes)
  575. class DenseNet169(BaseClassifier):
  576. def __init__(self, num_classes=1000):
  577. super(DenseNet169, self).__init__(
  578. model_name='DenseNet169', num_classes=num_classes)
  579. class DenseNet201(BaseClassifier):
  580. def __init__(self, num_classes=1000):
  581. super(DenseNet201, self).__init__(
  582. model_name='DenseNet201', num_classes=num_classes)
  583. class DenseNet264(BaseClassifier):
  584. def __init__(self, num_classes=1000):
  585. super(DenseNet264, self).__init__(
  586. model_name='DenseNet264', num_classes=num_classes)
  587. class HRNet_W18_C(BaseClassifier):
  588. def __init__(self, num_classes=1000):
  589. super(HRNet_W18_C, self).__init__(
  590. model_name='HRNet_W18_C', num_classes=num_classes)
  591. class HRNet_W30_C(BaseClassifier):
  592. def __init__(self, num_classes=1000):
  593. super(HRNet_W30_C, self).__init__(
  594. model_name='HRNet_W30_C', num_classes=num_classes)
  595. class HRNet_W32_C(BaseClassifier):
  596. def __init__(self, num_classes=1000):
  597. super(HRNet_W32_C, self).__init__(
  598. model_name='HRNet_W32_C', num_classes=num_classes)
  599. class HRNet_W40_C(BaseClassifier):
  600. def __init__(self, num_classes=1000):
  601. super(HRNet_W40_C, self).__init__(
  602. model_name='HRNet_W40_C', num_classes=num_classes)
  603. class HRNet_W44_C(BaseClassifier):
  604. def __init__(self, num_classes=1000):
  605. super(HRNet_W44_C, self).__init__(
  606. model_name='HRNet_W44_C', num_classes=num_classes)
  607. class HRNet_W48_C(BaseClassifier):
  608. def __init__(self, num_classes=1000):
  609. super(HRNet_W48_C, self).__init__(
  610. model_name='HRNet_W48_C', num_classes=num_classes)
  611. class HRNet_W64_C(BaseClassifier):
  612. def __init__(self, num_classes=1000):
  613. super(HRNet_W64_C, self).__init__(
  614. model_name='HRNet_W64_C', num_classes=num_classes)
  615. class Xception41(BaseClassifier):
  616. def __init__(self, num_classes=1000):
  617. super(Xception41, self).__init__(
  618. model_name='Xception41', num_classes=num_classes)
  619. class Xception65(BaseClassifier):
  620. def __init__(self, num_classes=1000):
  621. super(Xception65, self).__init__(
  622. model_name='Xception65', num_classes=num_classes)
  623. class Xception71(BaseClassifier):
  624. def __init__(self, num_classes=1000):
  625. super(Xception71, self).__init__(
  626. model_name='Xception71', num_classes=num_classes)
  627. class ShuffleNetV2(BaseClassifier):
  628. def __init__(self, num_classes=1000, scale=1.0):
  629. supported_scale = [.25, .33, .5, 1.0, 1.5, 2.0]
  630. if scale not in supported_scale:
  631. logging.warning("scale={} is not supported by ShuffleNetV2, "
  632. "scale is forcibly set to 1.0".format(scale))
  633. scale = 1.0
  634. model_name = 'ShuffleNetV2_x' + str(float(scale)).replace('.', '_')
  635. super(ShuffleNetV2, self).__init__(
  636. model_name=model_name, num_classes=num_classes)
  637. def _get_test_inputs(self, image_shape):
  638. if image_shape is not None:
  639. if len(image_shape) == 2:
  640. image_shape = [None, 3] + image_shape
  641. else:
  642. image_shape = [None, 3, 224, 224]
  643. logging.warning(
  644. '[Important!!!] When exporting inference model for {},'.format(
  645. self.__class__.__name__) +
  646. ' if fixed_input_shape is not set, it will be forcibly set to [None, 3, 224, 224]'
  647. +
  648. 'Please check image shape after transforms is [3, 224, 224], if not, fixed_input_shape '
  649. + 'should be specified manually.')
  650. self._fix_transforms_shape(image_shape[-2:])
  651. self.fixed_input_shape = image_shape
  652. input_spec = [
  653. InputSpec(
  654. shape=image_shape, name='image', dtype='float32')
  655. ]
  656. return input_spec
  657. class ShuffleNetV2_swish(BaseClassifier):
  658. def __init__(self, num_classes=1000):
  659. super(ShuffleNetV2_swish, self).__init__(
  660. model_name='ShuffleNetV2_x1_5', num_classes=num_classes)
  661. def _get_test_inputs(self, image_shape):
  662. if image_shape is not None:
  663. if len(image_shape) == 2:
  664. image_shape = [None, 3] + image_shape
  665. else:
  666. image_shape = [None, 3, 224, 224]
  667. logging.warning(
  668. '[Important!!!] When exporting inference model for {},'.format(
  669. self.__class__.__name__) +
  670. ' if fixed_input_shape is not set, it will be forcibly set to [None, 3, 224, 224]'
  671. +
  672. 'Please check image shape after transforms is [3, 224, 224], if not, fixed_input_shape '
  673. + 'should be specified manually.')
  674. self._fix_transforms_shape(image_shape[-2:])
  675. self.fixed_input_shape = image_shape
  676. input_spec = [
  677. InputSpec(
  678. shape=image_shape, name='image', dtype='float32')
  679. ]
  680. return input_spec