is_training = targets is not None
for yolo_layer in self.yolo_layers:
yolo_layer.img_dim = x.shape[2]
output = []
loss = 0
layer_outputs = []
for i, (module_def, module) in enumerate(zip(self.module_defs, self.module_list)):
if module_def["type"] in ["convolutional", "upsample", "maxpool"]:
x = module(x)
elif module_def["type"] == "route":
layer_i = [int(x) for x in module_def["layers"].split(",")]
x = torch.cat([layer_outputs[i] for i in layer_i], 1)
elif module_def["type"] == "shortcut":
layer_i = int(module_def["from"])
x = layer_outputs[-1] + layer_outputs[layer_i]
elif module_def["type"] == "yolo":
if is_training:
x, layer_loss = module[0](x, targets)
loss += layer_loss
else:
x = module(x)
output.append(x)
layer_outputs.append(x)
if is_training:
return loss, to_cpu(torch.cat(output, 1))
else:
return to_cpu(torch.cat(output, 1))
def load_darknet_weights(self, weights_path):
Parses and loads the weights stored in "weights_path"
