glenn-jocher commited on
Commit
e8cf24b
·
1 Parent(s): 121d90b

Initial model ensemble capability #318

Browse files
Files changed (4) hide show
  1. detect.py +3 -4
  2. models/experimental.py +17 -0
  3. test.py +12 -15
  4. utils/datasets.py +1 -1
detect.py CHANGED
@@ -2,7 +2,7 @@ import argparse
2
 
3
  import torch.backends.cudnn as cudnn
4
 
5
- from utils import google_utils
6
  from utils.datasets import *
7
  from utils.utils import *
8
 
@@ -20,8 +20,7 @@ def detect(save_img=False):
20
  half = device.type != 'cpu' # half precision only supported on CUDA
21
 
22
  # Load model
23
- google_utils.attempt_download(weights)
24
- model = torch.load(weights, map_location=device)['model'].float().eval() # load FP32 model
25
  imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
26
  if half:
27
  model.half() # to FP16
@@ -137,7 +136,7 @@ def detect(save_img=False):
137
 
138
  if __name__ == '__main__':
139
  parser = argparse.ArgumentParser()
140
- parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path')
141
  parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
142
  parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
143
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
 
2
 
3
  import torch.backends.cudnn as cudnn
4
 
5
+ from models.experimental import *
6
  from utils.datasets import *
7
  from utils.utils import *
8
 
 
20
  half = device.type != 'cpu' # half precision only supported on CUDA
21
 
22
  # Load model
23
+ model = attempt_load(weights, map_location=device) # load FP32 model
 
24
  imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
25
  if half:
26
  model.half() # to FP16
 
136
 
137
  if __name__ == '__main__':
138
  parser = argparse.ArgumentParser()
139
+ parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
140
  parser.add_argument('--source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam
141
  parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
142
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
models/experimental.py CHANGED
@@ -1,6 +1,7 @@
1
  # This file contains experimental modules
2
 
3
  from models.common import *
 
4
 
5
 
6
  class CrossConv(nn.Module):
@@ -119,3 +120,19 @@ class Ensemble(nn.ModuleList):
119
  for module in self:
120
  y.append(module(x, augment)[0])
121
  return torch.cat(y, 1), None # ensembled inference output, train output
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # This file contains experimental modules
2
 
3
  from models.common import *
4
+ from utils import google_utils
5
 
6
 
7
  class CrossConv(nn.Module):
 
120
  for module in self:
121
  y.append(module(x, augment)[0])
122
  return torch.cat(y, 1), None # ensembled inference output, train output
123
+
124
+
125
+ def attempt_load(weights, map_location=None):
126
+ # Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
127
+ model = Ensemble()
128
+ for w in weights if isinstance(weights, list) else [weights]:
129
+ google_utils.attempt_download(w)
130
+ model.append(torch.load(w, map_location=map_location)['model'].float().fuse().eval()) # load FP32 model
131
+
132
+ if len(model) == 1:
133
+ return model[-1] # return model
134
+ else:
135
+ print('Ensemble created with %s\n' % weights)
136
+ for k in ['names', 'stride']:
137
+ setattr(model, k, getattr(model[-1], k))
138
+ return model # return ensemble
test.py CHANGED
@@ -1,9 +1,8 @@
1
  import argparse
2
  import json
3
 
4
- from utils import google_utils
5
  from utils.datasets import *
6
- from utils.utils import *
7
 
8
 
9
  def test(data,
@@ -20,28 +19,26 @@ def test(data,
20
  dataloader=None,
21
  merge=False):
22
  # Initialize/load model and set device
23
- if model is None:
24
- training = False
25
- merge = opt.merge # use Merge NMS
 
 
26
  device = torch_utils.select_device(opt.device, batch_size=batch_size)
 
27
 
28
  # Remove previous
29
  for f in glob.glob('test_batch*.jpg'):
30
  os.remove(f)
31
 
32
  # Load model
33
- google_utils.attempt_download(weights)
34
- model = torch.load(weights, map_location=device)['model'].float().fuse().to(device) # load to FP32
35
  imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
36
 
37
  # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
38
  # if device.type != 'cpu' and torch.cuda.device_count() > 1:
39
  # model = nn.DataParallel(model)
40
 
41
- else: # called by train.py
42
- training = True
43
- device = next(model.parameters()).device # get model device
44
-
45
  # Half
46
  half = device.type != 'cpu' and torch.cuda.device_count() == 1 # half precision only supported on single-GPU
47
  if half:
@@ -56,11 +53,11 @@ def test(data,
56
  niou = iouv.numel()
57
 
58
  # Dataloader
59
- if dataloader is None: # not training
60
  img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
61
  _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
62
  path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
63
- dataloader = create_dataloader(path, imgsz, batch_size, int(max(model.stride)), opt,
64
  hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
65
 
66
  seen = 0
@@ -193,7 +190,7 @@ def test(data,
193
  if save_json and map50 and len(jdict):
194
  imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
195
  f = 'detections_val2017_%s_results.json' % \
196
- (weights.split(os.sep)[-1].replace('.pt', '') if weights else '') # filename
197
  print('\nCOCO mAP with pycocotools... saving %s...' % f)
198
  with open(f, 'w') as file:
199
  json.dump(jdict, file)
@@ -226,7 +223,7 @@ def test(data,
226
 
227
  if __name__ == '__main__':
228
  parser = argparse.ArgumentParser(prog='test.py')
229
- parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='model.pt path')
230
  parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
231
  parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
232
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
 
1
  import argparse
2
  import json
3
 
4
+ from models.experimental import *
5
  from utils.datasets import *
 
6
 
7
 
8
  def test(data,
 
19
  dataloader=None,
20
  merge=False):
21
  # Initialize/load model and set device
22
+ training = model is not None
23
+ if training: # called by train.py
24
+ device = next(model.parameters()).device # get model device
25
+
26
+ else: # called directly
27
  device = torch_utils.select_device(opt.device, batch_size=batch_size)
28
+ merge = opt.merge # use Merge NMS
29
 
30
  # Remove previous
31
  for f in glob.glob('test_batch*.jpg'):
32
  os.remove(f)
33
 
34
  # Load model
35
+ model = attempt_load(weights, map_location=device) # load FP32 model
 
36
  imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
37
 
38
  # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
39
  # if device.type != 'cpu' and torch.cuda.device_count() > 1:
40
  # model = nn.DataParallel(model)
41
 
 
 
 
 
42
  # Half
43
  half = device.type != 'cpu' and torch.cuda.device_count() == 1 # half precision only supported on single-GPU
44
  if half:
 
53
  niou = iouv.numel()
54
 
55
  # Dataloader
56
+ if not training:
57
  img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
58
  _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once
59
  path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
60
+ dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt,
61
  hyp=None, augment=False, cache=False, pad=0.5, rect=True)[0]
62
 
63
  seen = 0
 
190
  if save_json and map50 and len(jdict):
191
  imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files]
192
  f = 'detections_val2017_%s_results.json' % \
193
+ (weights.split(os.sep)[-1].replace('.pt', '') if isinstance(weights, str) else '') # filename
194
  print('\nCOCO mAP with pycocotools... saving %s...' % f)
195
  with open(f, 'w') as file:
196
  json.dump(jdict, file)
 
223
 
224
  if __name__ == '__main__':
225
  parser = argparse.ArgumentParser(prog='test.py')
226
+ parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
227
  parser.add_argument('--data', type=str, default='data/coco128.yaml', help='*.data path')
228
  parser.add_argument('--batch-size', type=int, default=32, help='size of each image batch')
229
  parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
utils/datasets.py CHANGED
@@ -48,7 +48,7 @@ def create_dataloader(path, imgsz, batch_size, stride, opt, hyp=None, augment=Fa
48
  rect=rect, # rectangular training
49
  cache_images=cache,
50
  single_cls=opt.single_cls,
51
- stride=stride,
52
  pad=pad)
53
 
54
  batch_size = min(batch_size, len(dataset))
 
48
  rect=rect, # rectangular training
49
  cache_images=cache,
50
  single_cls=opt.single_cls,
51
+ stride=int(stride),
52
  pad=pad)
53
 
54
  batch_size = min(batch_size, len(dataset))