Ankle/yolov5_birdproject
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Add verbose option to pytorch hub models (#2926)

Browse Source 
* Add verbose and update print to logging

* Fix positonal param

* Revert auto formatting changes

* Update hubconf.py

Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
NanoCode012andGlenn Jocher authored and GitHub committed 2021-04-25 22:01:05 +02:00
1 parent 3665c0f59b
commit c0d3f80544
3 files changed
+41 -38
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

hubconf.py
+28 -25
View File
@@ -16,10 +16,9 @@ from utils.torch_utils import select_device
dependencies = ['torch', 'yaml'] dependencies = ['torch', 'yaml']
check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop')) check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('pycocotools', 'thop'))
set_logging()
def create(name, pretrained, channels, classes, autoshape): def create(name, pretrained, channels, classes, autoshape, verbose):
"""Creates a specified YOLOv5 model """Creates a specified YOLOv5 model
Arguments: Arguments:
@@ -32,6 +31,8 @@ def create(name, pretrained, channels, classes, autoshape):
pytorch model pytorch model
""" """
try: try:
set_logging(verbose=verbose)
cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path
model = Model(cfg, channels, classes) model = Model(cfg, channels, classes)
if pretrained: if pretrained:
@@ -55,7 +56,7 @@ def create(name, pretrained, channels, classes, autoshape):
raise Exception(s) from e raise Exception(s) from e
def custom(path_or_model='path/to/model.pt', autoshape=True): def custom(path_or_model='path/to/model.pt', autoshape=True, verbose=True):
"""YOLOv5-custom model https://github.com/ultralytics/yolov5 """YOLOv5-custom model https://github.com/ultralytics/yolov5
Arguments (3 options): Arguments (3 options):
@@ -66,6 +67,8 @@ def custom(path_or_model='path/to/model.pt', autoshape=True):
Returns: Returns:
pytorch model pytorch model
""" """
set_logging(verbose=verbose)
model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
if isinstance(model, dict): if isinstance(model, dict):
model = model['ema' if model.get('ema') else 'model'] # load model model = model['ema' if model.get('ema') else 'model'] # load model
@@ -79,49 +82,49 @@ def custom(path_or_model='path/to/model.pt', autoshape=True):
return hub_model.to(device) return hub_model.to(device)
def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-small model https://github.com/ultralytics/yolov5 # YOLOv5-small model https://github.com/ultralytics/yolov5
return create('yolov5s', pretrained, channels, classes, autoshape) return create('yolov5s', pretrained, channels, classes, autoshape, verbose)
def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-medium model https://github.com/ultralytics/yolov5 # YOLOv5-medium model https://github.com/ultralytics/yolov5
return create('yolov5m', pretrained, channels, classes, autoshape) return create('yolov5m', pretrained, channels, classes, autoshape, verbose)
def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-large model https://github.com/ultralytics/yolov5 # YOLOv5-large model https://github.com/ultralytics/yolov5
return create('yolov5l', pretrained, channels, classes, autoshape) return create('yolov5l', pretrained, channels, classes, autoshape, verbose)
def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-xlarge model https://github.com/ultralytics/yolov5 # YOLOv5-xlarge model https://github.com/ultralytics/yolov5
return create('yolov5x', pretrained, channels, classes, autoshape) return create('yolov5x', pretrained, channels, classes, autoshape, verbose)
def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-small model https://github.com/ultralytics/yolov5 # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5
return create('yolov5s6', pretrained, channels, classes, autoshape) return create('yolov5s6', pretrained, channels, classes, autoshape, verbose)
def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-medium model https://github.com/ultralytics/yolov5 # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5
return create('yolov5m6', pretrained, channels, classes, autoshape) return create('yolov5m6', pretrained, channels, classes, autoshape, verbose)
def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-large model https://github.com/ultralytics/yolov5 # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5
return create('yolov5l6', pretrained, channels, classes, autoshape) return create('yolov5l6', pretrained, channels, classes, autoshape, verbose)
def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True): def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True):
# YOLOv5-xlarge model https://github.com/ultralytics/yolov5 # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5
return create('yolov5x6', pretrained, channels, classes, autoshape) return create('yolov5x6', pretrained, channels, classes, autoshape, verbose)
if __name__ == '__main__': if __name__ == '__main__':
model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True) # pretrained example model = create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained
# model = custom(path_or_model='path/to/model.pt') # custom example # model = custom(path_or_model='path/to/model.pt') # custom
# Verify inference # Verify inference
import cv2 import cv2
models/yolo.py
+11 -11
View File
@@ -84,7 +84,7 @@ class Model(nn.Module):
self.yaml['anchors'] = round(anchors) # override yaml value self.yaml['anchors'] = round(anchors) # override yaml value
self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist
self.names = [str(i) for i in range(self.yaml['nc'])] # default names self.names = [str(i) for i in range(self.yaml['nc'])] # default names
# print([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))]) # logger.info([x.shape for x in self.forward(torch.zeros(1, ch, 64, 64))])
# Build strides, anchors # Build strides, anchors
m = self.model[-1] # Detect() m = self.model[-1] # Detect()
@@ -95,7 +95,7 @@ class Model(nn.Module):
check_anchor_order(m) check_anchor_order(m)
self.stride = m.stride self.stride = m.stride
self._initialize_biases() # only run once self._initialize_biases() # only run once
# print('Strides: %s' % m.stride.tolist()) # logger.info('Strides: %s' % m.stride.tolist())
# Init weights, biases # Init weights, biases
initialize_weights(self) initialize_weights(self)
@@ -134,13 +134,13 @@ class Model(nn.Module):
for _ in range(10): for _ in range(10):
_ = m(x) _ = m(x)
dt.append((time_synchronized() - t) * 100) dt.append((time_synchronized() - t) * 100)
print('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type)) logger.info('%10.1f%10.0f%10.1fms %-40s' % (o, m.np, dt[-1], m.type))
x = m(x) # run x = m(x) # run
y.append(x if m.i in self.save else None) # save output y.append(x if m.i in self.save else None) # save output
if profile: if profile:
print('%.1fms total' % sum(dt)) logger.info('%.1fms total' % sum(dt))
return x return x
def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency def _initialize_biases(self, cf=None): # initialize biases into Detect(), cf is class frequency
@@ -157,15 +157,15 @@ class Model(nn.Module):
m = self.model[-1] # Detect() module m = self.model[-1] # Detect() module
for mi in m.m: # from for mi in m.m: # from
b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85)
print(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) logger.info(('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()))
# def _print_weights(self): # def _print_weights(self):
# for m in self.model.modules(): # for m in self.model.modules():
# if type(m) is Bottleneck: # if type(m) is Bottleneck:
# print('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights # logger.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
print('Fusing layers... ') logger.info('Fusing layers... ')
for m in self.model.modules(): for m in self.model.modules():
if type(m) is Conv and hasattr(m, 'bn'): if type(m) is Conv and hasattr(m, 'bn'):
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
@@ -177,19 +177,19 @@ class Model(nn.Module):
def nms(self, mode=True): # add or remove NMS module def nms(self, mode=True): # add or remove NMS module
present = type(self.model[-1]) is NMS # last layer is NMS present = type(self.model[-1]) is NMS # last layer is NMS
if mode and not present: if mode and not present:
print('Adding NMS... ') logger.info('Adding NMS... ')
m = NMS() # module m = NMS() # module
m.f = -1 # from m.f = -1 # from
m.i = self.model[-1].i + 1 # index m.i = self.model[-1].i + 1 # index
self.model.add_module(name='%s' % m.i, module=m) # add self.model.add_module(name='%s' % m.i, module=m) # add
self.eval() self.eval()
elif not mode and present: elif not mode and present:
print('Removing NMS... ') logger.info('Removing NMS... ')
self.model = self.model[:-1] # remove self.model = self.model[:-1] # remove
return self return self
def autoshape(self): # add autoShape module def autoshape(self): # add autoShape module
print('Adding autoShape... ') logger.info('Adding autoShape... ')
m = autoShape(self) # wrap model m = autoShape(self) # wrap model
copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes
return m return m
@@ -272,6 +272,6 @@ if __name__ == '__main__':
# Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
# from torch.utils.tensorboard import SummaryWriter # from torch.utils.tensorboard import SummaryWriter
# tb_writer = SummaryWriter('.') # tb_writer = SummaryWriter('.')
# print("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/") # logger.info("Run 'tensorboard --logdir=models' to view tensorboard at http://localhost:6006/")
# tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph # tb_writer.add_graph(torch.jit.trace(model, img, strict=False), []) # add model graph
# tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard # tb_writer.add_image('test', img[0], dataformats='CWH') # add model to tensorboard
utils/general.py
+2 -2
View File
@@ -32,10 +32,10 @@ cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with Py
os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads
def set_logging(rank=-1): def set_logging(rank=-1, verbose=True):
logging.basicConfig( logging.basicConfig(
format="%(message)s", format="%(message)s",
level=logging.INFO if rank in [-1, 0] else logging.WARN) level=logging.INFO if (verbose and rank in [-1, 0]) else logging.WARN)
def init_seeds(seed=0): def init_seeds(seed=0):