Register
Login
Resources
Docs Blog Datasets Glossary Case Studies Tutorials & Webinars
Product
Data Engine LLMs Platform Enterprise
Pricing Explore
Connect to our Discord channel

pretrained_models.py 6.5 KB

You have to be logged in to leave a comment. Sign In
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
  1. # TODO: It would be nice to create keys here as: make_pretrained_model_key(Models.RESNET18, Dataset.COCO)
  2. # TODO: Not only this would reduce risk of making a typo error, it would bring more clarity how the key is created
  3. # TODO: And allow to "query" pretrained models by dataset
  4. MODEL_URLS = {
  5. "regnetY800_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/regnetY800_imagenet.pth",
  6. "regnetY600_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/regnetY600_imagenet.pth",
  7. "regnetY400_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/regnetY400_imagenet.pth",
  8. "regnetY200_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/regnetY200_imagenet.pth",
  9. "resnet50_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/resnet50_imagenet.pth",
  10. "resnet34_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/resnet34_imagenet.pth",
  11. "resnet18_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/resnet18_imagenet.pth",
  12. "repvgg_a0_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/repvgg_a0_imagenet.pth",
  13. "shelfnet34_lw_coco_segmentation_subclass": "https://sg-hub-nv.s3.amazonaws.com/models/shelfnet34_lw_coco_segmentation_subclass.pth",
  14. "ddrnet_23_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/ddrnet_23_cityscapes.pth",
  15. "ddrnet_23_slim_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/ddrnet_23_slim_cityscapes.pth",
  16. "ddrnet_39_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/ddrnet_39_cityscapes.pth",
  17. "stdc1_seg50_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/stdc1_seg50_cityscapes.pth",
  18. "stdc1_seg75_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/stdc1_seg75_cityscapes.pth",
  19. "stdc2_seg50_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/stdc2_seg50_cityscapes.pth",
  20. "stdc2_seg75_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/stdc2_seg75_cityscapes.pth",
  21. "efficientnet_b0_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/efficientnet_b0_imagenet.pth",
  22. "ssd_lite_mobilenet_v2_coco": "https://sg-hub-nv.s3.amazonaws.com/models/ssd_lite_mobilenet_v2_coco.pth",
  23. "ssd_mobilenet_v1_coco": "https://sg-hub-nv.s3.amazonaws.com/models/ssd_mobilenet_v1_coco.pth",
  24. "mobilenet_v3_large_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/mobilenet_v3_large_imagenet.pth",
  25. "mobilenet_v3_small_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/mobilenet_v3_small_imagenet.pth",
  26. "mobilenet_v2_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/mobilenet_v2_imagenet.pth",
  27. "regseg48_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/regseg48_cityscapes.pth",
  28. "vit_base_imagenet21k": "https://sg-hub-nv.s3.amazonaws.com/models/vit_base_imagenet21k.pth",
  29. "vit_large_imagenet21k": "https://sg-hub-nv.s3.amazonaws.com/models/vit_large_imagenet21k.pth",
  30. "vit_base_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/vit_base_imagenet.pth",
  31. "vit_large_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/vit_large_imagenet.pth",
  32. "beit_base_patch16_224_imagenet": "https://sg-hub-nv.s3.amazonaws.com/models/beit_base_patch16_224_imagenet.pth",
  33. "beit_base_patch16_224_cifar10": "https://sg-hub-nv.s3.amazonaws.com/models/beit_base_patch16_224_cifar10.pth",
  34. "yolox_s_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolox_s_coco.pth",
  35. "yolox_m_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolox_m_coco.pth",
  36. "yolox_l_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolox_l_coco.pth",
  37. "yolox_t_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolox_t_coco.pth",
  38. "yolox_n_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolox_n_coco.pth",
  39. "pp_lite_t_seg50_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/pp_lite_t_seg50_cityscapes.pth",
  40. "pp_lite_t_seg75_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/pp_lite_t_seg75_cityscapes.pth",
  41. "pp_lite_b_seg50_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/pp_lite_b_seg50_cityscapes.pth",
  42. "pp_lite_b_seg75_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/pp_lite_b_seg75_cityscapes.pth",
  43. "ppyoloe_s_coco": "https://sg-hub-nv.s3.amazonaws.com/models/ppyoloe_s_coco.pth",
  44. "ppyoloe_m_coco": "https://sg-hub-nv.s3.amazonaws.com/models/ppyoloe_m_coco.pth",
  45. "ppyoloe_l_coco": "https://sg-hub-nv.s3.amazonaws.com/models/ppyoloe_l_coco.pth",
  46. "ppyoloe_x_coco": "https://sg-hub-nv.s3.amazonaws.com/models/ppyoloe_x_coco.pth",
  47. "yolo_nas_s_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_s_coco.pth",
  48. "yolo_nas_m_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_m_coco.pth",
  49. "yolo_nas_l_coco": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_l_coco.pth",
  50. "dekr_w32_no_dc_coco_pose": "https://sg-hub-nv.s3.amazonaws.com/models/dekr_w32_no_dc_coco_pose.pth",
  51. "pose_rescoring_coco_coco_pose": "https://sg-hub-nv.s3.amazonaws.com/models/pose_rescoring_coco_coco_pose.pth",
  52. "segformer_b0_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/segformer_b0_cityscapes.pth",
  53. "segformer_b1_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/segformer_b1_cityscapes.pth",
  54. "segformer_b2_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/segformer_b2_cityscapes.pth",
  55. "segformer_b3_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/segformer_b3_cityscapes.pth",
  56. "segformer_b4_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/segformer_b4_cityscapes.pth",
  57. "segformer_b5_cityscapes": "https://sg-hub-nv.s3.amazonaws.com/models/segformer_b5_cityscapes.pth",
  58. "yolo_nas_pose_n_coco_pose": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_pose_n_coco_pose.pth",
  59. "yolo_nas_pose_s_coco_pose": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_pose_s_coco_pose.pth",
  60. "yolo_nas_pose_m_coco_pose": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_pose_m_coco_pose.pth",
  61. "yolo_nas_pose_l_coco_pose": "https://sg-hub-nv.s3.amazonaws.com/models/yolo_nas_pose_l_coco_pose.pth",
  62. }
  63. PRETRAINED_NUM_CLASSES = {
  64. "imagenet": 1000,
  65. "imagenet21k": 21843,
  66. "coco_segmentation_subclass": 21,
  67. "cityscapes": 19,
  68. "coco": 80,
  69. "coco_pose": 17,
  70. "cifar10": 10,
  71. }
  72. DATASET_LICENSES = {
  73. "imagenet": "https://www.image-net.org/download.php",
  74. "imagenet21k": "https://github.com/Alibaba-MIIL/ImageNet21K",
  75. "coco": "https://cocodataset.org/#termsofuse",
  76. "coco_segmentation_subclass": "https://cocodataset.org/#termsofuse",
  77. "coco_pose": "https://cocodataset.org/#termsofuse",
  78. "cityscapes": "https://www.cs.toronto.edu/~kriz/cifar.html",
  79. "objects365": "https://www.objects365.org/download.html",
  80. }
Tip!

Press p or to see the previous file or, n or to see the next file

Comments

Loading...