Browse Source

Remove `.autoshape()` method (#5694)

modifyDataloader
Glenn Jocher GitHub 3 years ago
parent
commit
46daa7b78d
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 6 additions and 14 deletions
  1. +2
    -1
      hubconf.py
  2. +3
    -5
      models/common.py
  3. +1
    -8
      models/yolo.py

+ 2
- 1
hubconf.py View File

""" """
from pathlib import Path from pathlib import Path


from models.common import AutoShape
from models.experimental import attempt_load from models.experimental import attempt_load
from models.yolo import Model from models.yolo import Model
from utils.downloads import attempt_download from utils.downloads import attempt_download
if len(ckpt['model'].names) == classes: if len(ckpt['model'].names) == classes:
model.names = ckpt['model'].names # set class names attribute model.names = ckpt['model'].names # set class names attribute
if autoshape: if autoshape:
model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS
model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS
return model.to(device) return model.to(device)


except Exception as e: except Exception as e:

+ 3
- 5
models/common.py View File

from utils.general import (LOGGER, check_requirements, check_suffix, colorstr, increment_path, make_divisible, from utils.general import (LOGGER, check_requirements, check_suffix, colorstr, increment_path, make_divisible,
non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh) non_max_suppression, scale_coords, xywh2xyxy, xyxy2xywh)
from utils.plots import Annotator, colors, save_one_box from utils.plots import Annotator, colors, save_one_box
from utils.torch_utils import time_sync
from utils.torch_utils import copy_attr, time_sync




def autopad(k, p=None): # kernel, padding def autopad(k, p=None): # kernel, padding


def __init__(self, model): def __init__(self, model):
super().__init__() super().__init__()
LOGGER.info('Adding AutoShape... ')
copy_attr(self, model, include=('yaml', 'nc', 'hyp', 'names', 'stride', 'abc'), exclude=()) # copy attributes
self.model = model.eval() self.model = model.eval()


def autoshape(self):
LOGGER.info('AutoShape already enabled, skipping... ') # model already converted to model.autoshape()
return self

def _apply(self, fn): def _apply(self, fn):
# Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers # Apply to(), cpu(), cuda(), half() to model tensors that are not parameters or registered buffers
self = super()._apply(fn) self = super()._apply(fn)

+ 1
- 8
models/yolo.py View File

from utils.autoanchor import check_anchor_order from utils.autoanchor import check_anchor_order
from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args
from utils.plots import feature_visualization from utils.plots import feature_visualization
from utils.torch_utils import (copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device,
time_sync)
from utils.torch_utils import fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device, time_sync


try: try:
import thop # for FLOPs computation import thop # for FLOPs computation
self.info() self.info()
return self return self


def autoshape(self): # add AutoShape module
LOGGER.info('Adding AutoShape... ')
m = AutoShape(self) # wrap model
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
return m

def info(self, verbose=False, img_size=640): # print model information def info(self, verbose=False, img_size=640): # print model information
model_info(self, verbose, img_size) model_info(self, verbose, img_size)



Loading…
Cancel
Save