diff --git a/models.py b/models.py index eb2fdcea..119aaf4b 100755 --- a/models.py +++ b/models.py @@ -96,7 +96,7 @@ class YOLOLayer(nn.Module): def __init__(self, anchors, nc, img_size, yolo_index): super(YOLOLayer, self).__init__() - self.anchors = torch.Tensor(anchors) + self.anchors = torch.from_numpy(anchors) self.na = len(anchors) # number of anchors (3) self.nc = nc # number of classes (80) self.nx = 0 # initialize number of x gridpoints diff --git a/train.py b/train.py index 0166dccb..c0103aa2 100644 --- a/train.py +++ b/train.py @@ -200,8 +200,7 @@ def train(cfg, # Start training model.nc = nc # attach number of classes to model model.hyp = hyp # attach hyperparameters to model - if dataset.image_weights: - model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights + model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights model_info(model, report='summary') # 'full' or 'summary' nb = len(dataloader) maps = np.zeros(nc) # mAP per class diff --git a/utils/utils.py b/utils/utils.py index 0a0ed43c..497dab38 100755 --- a/utils/utils.py +++ b/utils/utils.py @@ -62,7 +62,7 @@ def labels_to_class_weights(labels, nc=80): weights[weights == 0] = 1 # replace empty bins with 1 weights = 1 / weights # number of targets per class weights /= weights.sum() # normalize - return torch.Tensor(weights) + return torch.from_numpy(weights) def labels_to_image_weights(labels, nc=80, class_weights=np.ones(80)):