This commit is contained in:
Glenn Jocher 2020-03-09 14:20:38 -07:00
parent 5fb661b7d4
commit 67e7ac221f
1 changed files with 2 additions and 3 deletions

View File

@ -195,7 +195,7 @@ class YOLOLayer(nn.Module):
ny = img_size[0] // stride # number y grid points ny = img_size[0] // stride # number y grid points
create_grids(self, img_size, (nx, ny)) create_grids(self, img_size, (nx, ny))
def forward(self, p, img_size, var=None): def forward(self, p, img_size):
if ONNX_EXPORT: if ONNX_EXPORT:
bs = 1 # batch size bs = 1 # batch size
else: else:
@ -261,10 +261,9 @@ class Darknet(nn.Module):
self.version = np.array([0, 2, 5], dtype=np.int32) # (int32) version info: major, minor, revision self.version = np.array([0, 2, 5], dtype=np.int32) # (int32) version info: major, minor, revision
self.seen = np.array([0], dtype=np.int64) # (int64) number of images seen during training self.seen = np.array([0], dtype=np.int64) # (int64) number of images seen during training
def forward(self, x, var=None): def forward(self, x, verbose=False):
img_size = x.shape[-2:] img_size = x.shape[-2:]
yolo_out, out = [], [] yolo_out, out = [], []
verbose = False
if verbose: if verbose:
str = '' str = ''
print('0', x.shape) print('0', x.shape)