|
|
@@ -265,21 +265,13 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary |
|
|
|
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ |
|
|
|
model.train() |
|
|
|
|
|
|
|
# Update image weights (optional) |
|
|
|
# Update image weights (optional, single-GPU only) |
|
|
|
if opt.image_weights: |
|
|
|
# Generate indices |
|
|
|
if RANK in [-1, 0]: |
|
|
|
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights |
|
|
|
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights |
|
|
|
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx |
|
|
|
# Broadcast if DDP |
|
|
|
if RANK != -1: |
|
|
|
indices = (torch.tensor(dataset.indices) if RANK == 0 else torch.zeros(dataset.n)).int() |
|
|
|
dist.broadcast(indices, 0) |
|
|
|
if RANK != 0: |
|
|
|
dataset.indices = indices.cpu().numpy() |
|
|
|
|
|
|
|
# Update mosaic border |
|
|
|
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights |
|
|
|
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights |
|
|
|
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx |
|
|
|
|
|
|
|
# Update mosaic border (optional) |
|
|
|
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) |
|
|
|
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders |
|
|
|
|