|
|
@@ -237,8 +237,8 @@ def parse_model(d, ch): # model_dict, input_channels(3) |
|
|
|
pass |
|
|
|
|
|
|
|
n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain |
|
|
|
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, |
|
|
|
C3, C3TR, C3SPP, C3Ghost]: |
|
|
|
if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, |
|
|
|
BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]: |
|
|
|
c1, c2 = ch[f], args[0] |
|
|
|
if c2 != no: # if not output |
|
|
|
c2 = make_divisible(c2 * gw, 8) |
|
|
@@ -279,6 +279,7 @@ if __name__ == '__main__': |
|
|
|
parser = argparse.ArgumentParser() |
|
|
|
parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') |
|
|
|
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') |
|
|
|
parser.add_argument('--profile', action='store_true', help='profile model speed') |
|
|
|
opt = parser.parse_args() |
|
|
|
opt.cfg = check_file(opt.cfg) # check file |
|
|
|
set_logging() |
|
|
@@ -289,8 +290,9 @@ if __name__ == '__main__': |
|
|
|
model.train() |
|
|
|
|
|
|
|
# Profile |
|
|
|
# img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 320, 320).to(device) |
|
|
|
# y = model(img, profile=True) |
|
|
|
if opt.profile: |
|
|
|
img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device) |
|
|
|
y = model(img, profile=True) |
|
|
|
|
|
|
|
# Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) |
|
|
|
# from torch.utils.tensorboard import SummaryWriter |