@@ -121,7 +121,7 @@ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, nu
121121 def forward (self , x ):
122122 y1 = self .cv3 (self .m (self .cv1 (x )))
123123 y2 = self .cv2 (x )
124- return self .cv4 (self .act (self .bn (torch .cat ((y1 , y2 ), dim = 1 ))))
124+ return self .cv4 (self .act (self .bn (torch .cat ((y1 , y2 ), 1 ))))
125125
126126
127127class C3 (nn .Module ):
@@ -136,7 +136,7 @@ def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): # ch_in, ch_out, nu
136136 # self.m = nn.Sequential(*(CrossConv(c_, c_, 3, 1, g, 1.0, shortcut) for _ in range(n)))
137137
138138 def forward (self , x ):
139- return self .cv3 (torch .cat ((self .m (self .cv1 (x )), self .cv2 (x )), dim = 1 ))
139+ return self .cv3 (torch .cat ((self .m (self .cv1 (x )), self .cv2 (x )), 1 ))
140140
141141
142142class C3TR (C3 ):
@@ -527,7 +527,7 @@ def forward(self, imgs, size=640, augment=False, profile=False):
527527 p = next (self .model .parameters ()) if self .pt else torch .zeros (1 ) # for device and type
528528 autocast = self .amp and (p .device .type != 'cpu' ) # Automatic Mixed Precision (AMP) inference
529529 if isinstance (imgs , torch .Tensor ): # torch
530- with amp .autocast (enabled = autocast ):
530+ with amp .autocast (autocast ):
531531 return self .model (imgs .to (p .device ).type_as (p ), augment , profile ) # inference
532532
533533 # Pre-process
@@ -550,19 +550,19 @@ def forward(self, imgs, size=640, augment=False, profile=False):
550550 shape1 .append ([y * g for y in s ])
551551 imgs [i ] = im if im .data .contiguous else np .ascontiguousarray (im ) # update
552552 shape1 = [make_divisible (x , self .stride ) if self .pt else size for x in np .array (shape1 ).max (0 )] # inf shape
553- x = [letterbox (im , new_shape = shape1 , auto = False )[0 ] for im in imgs ] # pad
553+ x = [letterbox (im , shape1 , auto = False )[0 ] for im in imgs ] # pad
554554 x = np .ascontiguousarray (np .array (x ).transpose ((0 , 3 , 1 , 2 ))) # stack and BHWC to BCHW
555555 x = torch .from_numpy (x ).to (p .device ).type_as (p ) / 255 # uint8 to fp16/32
556556 t .append (time_sync ())
557557
558- with amp .autocast (enabled = autocast ):
558+ with amp .autocast (autocast ):
559559 # Inference
560560 y = self .model (x , augment , profile ) # forward
561561 t .append (time_sync ())
562562
563563 # Post-process
564- y = non_max_suppression (y if self .dmb else y [0 ], self .conf , iou_thres = self .iou , classes = self .classes ,
565- agnostic = self . agnostic , multi_label = self .multi_label , max_det = self .max_det ) # NMS
564+ y = non_max_suppression (y if self .dmb else y [0 ], self .conf , self .iou , self .classes , self . agnostic ,
565+ self .multi_label , max_det = self .max_det ) # NMS
566566 for i in range (n ):
567567 scale_coords (shape1 , y [i ][:, :4 ], shape0 [i ])
568568
0 commit comments