Some issue about the output of model in train.py and detect.py

[Louiszhu970814]

Hi I have some questions about the output prediction of the model

In train.py when model predicts a prediction

pred = model(imgs)
imgs is a tensor of [batch_size, 3, image_size, image_size] and pred is a list with len(pred)==3

while in detect.py the prediction code is this one

pred = model(img, augment=opt.augment)[0]
In this case, input is a tensor of [1, 3, image_size, image_size], pred is a tensor which is very different from the one in train.py

And if I modified the code to this pred = model(img, augment=opt.augment).
I delete "[0]"
and pred will become a tuple.

Does that mean the model in train.py is different from the model in detect.py?
Can anyone explain why pred of train.py looks like that??
Thank you guys!!

[glenn-jocher]

@Louiszhu970814 YOLOv5 output varies by train, eval mode type:

yolov5/models/yolo.py

Line 68 in 4695ca8

return x if self.training else (torch.cat(z, 1), x)