henry000 commited on
Commit
198ddce
·
1 Parent(s): 425a4ef

🎨 [Update] forward proccess, only save usable data

Browse files
Files changed (1) hide show
  1. yolo/model/yolo.py +4 -3
yolo/model/yolo.py CHANGED
@@ -58,15 +58,16 @@ class YOLO(nn.Module):
58
  layer_idx += 1
59
 
60
  def forward(self, x):
61
- y = [x]
62
  output = []
63
- for layer in self.model:
64
  if isinstance(layer.source, list):
65
  model_input = [y[idx] for idx in layer.source]
66
  else:
67
  model_input = y[layer.source]
68
  x = layer(model_input)
69
- y.append(x)
 
70
  if layer.output:
71
  output.append(x)
72
  return output
 
58
  layer_idx += 1
59
 
60
  def forward(self, x):
61
+ y = {0: x}
62
  output = []
63
+ for index, layer in enumerate(self.model, start=1):
64
  if isinstance(layer.source, list):
65
  model_input = [y[idx] for idx in layer.source]
66
  else:
67
  model_input = y[layer.source]
68
  x = layer(model_input)
69
+ if hasattr(layer, "save"):
70
+ y[index] = x
71
  if layer.output:
72
  output.append(x)
73
  return output