Lijun Yu
4 years ago
1 changed files with 5 additions and 1 deletions
+ 5
- 1
hubconf.py
View File
| """ | """ | ||||
| dependencies = ['torch', 'yaml'] | dependencies = ['torch', 'yaml'] | ||||
| import os | |||||
| import torch | import torch | ||||
| from models.yolo import Model | from models.yolo import Model | ||||
| Returns: | Returns: | ||||
| pytorch model | pytorch model | ||||
| """ | """ | ||||
| model = Model('models/%s.yaml' % name, channels, classes) | |||||
| config = os.path.join(os.path.dirname(__file__), 'models', '%s.yaml' % name) | |||||
| model = Model(config, channels, classes) | |||||
| if pretrained: | if pretrained: | ||||
| ckpt = '%s.pt' % name # checkpoint filename | ckpt = '%s.pt' % name # checkpoint filename | ||||
| google_utils.attempt_download(ckpt) # download if not found locally | google_utils.attempt_download(ckpt) # download if not found locally |
Loading…