XinzeLee / PolygonObjectDetection

This repository is based on Ultralytics/yolov5, with adjustments to enable polygon prediction boxes.
352 stars 95 forks source link

Expected all tensors to be on the same device, but found at least two devices, cuda: 0 and cpu! #30

Closed sac3tf closed 2 years ago

sac3tf commented 2 years ago

When following the first tutorial on Google Colab, I am trying to run !python polygon_test.py --weights polygon-yolov5s-ucas.pt --data polygon_ucas.yaml --img 1024 --iou 0.65 --task val --device 0 as in the example. I get the following error:

Traceback (most recent call last): File "polygon_test.py", line 325, in <module> test(**vars(opt)) File "/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py", line 28, in decorate_context return func(*args, **kwargs) File "polygon_test.py", line 224, in test for j in (ious > iouv[index_ap50]).nonzero(as_tuple=False): RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

I have not modified any code or data and cannot figure out where the issue is. Any help would be much appreciated. Thank you!

sac3tf commented 2 years ago

I was able to fix this and subsequent errors by using this script with several changes in the polygon_test.py file.

Two to note:

Line 153 needed to add .to('cpu') after each element in the stats.append (except tcls) Lines 224 and 229 had to add .to(device) after both ious and iouv

`import argparse import json import os from pathlib import Path from threading import Thread

import time import numpy as np import torch import yaml from tqdm import tqdm

from models.experimental import attempt_load from utils.datasets import create_dataloader # specify polygon=True to enable polygon anchor boxes from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, check_requirements, \ box_iou, non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path, colorstr, \ polygon_box_iou, polygon_non_max_suppression, polygon_scale_coords from utils.metrics import ap_per_class, ConfusionMatrix, Polygon_ConfusionMatrix from utils.plots import plot_images, output_to_target, plot_study_txt, polygon_plot_images, polygon_output_to_target from utils.torch_utils import select_device, time_synchronized

polygon_test.py is specially for polygon boxes

@torch.no_grad() def test(data, weights=None, # model.pt path(s) batch_size=32, # batch size imgsz=640, # inference size (pixels) conf_thres=0.001, # confidence threshold iou_thres=0.6, # NMS IoU threshold task='val', # train, val, test, speed or study device=0, # cuda device, i.e. 0 or 0,1,2,3 or cpu single_cls=False, # treat as single-class dataset augment=False, # augmented inference verbose=False, # verbose output save_txt=False, # save results to .txt save_hybrid=False, # save label+prediction hybrid results to .txt save_conf=False, # save confidences in --save-txt labels save_json=False, # save a cocoapi-compatible JSON results file project='runs/test', # save to project/name name='exp', # save to project/name exist_ok=False, # existing project/name ok, do not increment half=True, # use FP16 half-precision inference model=None, dataloader=None, save_dir=Path(''), plots=True, wandb_logger=None, compute_loss=None, ):

Initialize/load model and set device

training = model is not None
if training:  # called by polygon_train.py
    device = 0  # get model device

else:  # called directly
    set_logging()
    device = select_device(device, batch_size=batch_size)
    print(device)

    # Directories
    save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)  # increment run
    (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True)  # make dir

    # Load model
    model = attempt_load(weights, map_location=device)  # load FP32 model
    gs = max(int(model.stride.max()), 32)  # grid size (max stride)
    imgsz = check_img_size(imgsz, s=gs)  # check image size

    # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
    # if device.type != 'cpu' and torch.cuda.device_count() > 1:
    #     model = nn.DataParallel(model)

# Half
half &= device.type != 'cpu'  # half precision only supported on CUDA
if half:
    model.half().to(device)

# Configure
model.eval().to(device)
if isinstance(data, str):
    with open(data) as f:
        data = yaml.safe_load(f)
check_dataset(data)  # check
nc = 1 if single_cls else int(data['nc'])  # number of classes
is_coco = data['val'].endswith('coco/val2017.txt') and (nc==80)  # COCO dataset
iouv = torch.linspace(0.5, 0.95, 10).to(device)  # iou vector for mAP@0.5:0.95
index_ap50 = 0   # index for ap 0.5
niou = iouv.numel()

# Logging
log_imgs = 0
if wandb_logger and wandb_logger.wandb:
    log_imgs = min(wandb_logger.log_imgs, 100)
# Dataloader
if not training:
    if device.type != 'cpu':
        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once
    task = task if task in ('train', 'val', 'test') else 'val'  # path to train/val/test images
    dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=0.5, rect=True,
                                   prefix=colorstr(f'{task}: '), polygon=True)[0]    # POLYGON: polygon=True

seen = 0
confusion_matrix = Polygon_ConfusionMatrix(nc=nc)
names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)}
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
p, r, f1, mp, mr, map50, map, t0, t1, t2 = 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.
loss = torch.zeros(3, device=device)
jdict, stats, ap, ap_class = [], [], [], []
for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)):
    t_ = time_synchronized()
    img = img.to(device, non_blocking=True)
    img = img.half() if half else img.float()  # uint8 to fp16/32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    targets = targets.to(device)
    nb, _, height, width = img.shape  # batch size, channels, height, width
    t = time_synchronized()
    t0 += t - t_

    # Run model
    print("LOOK HERE!!!!!!")
    print("HEYYYY THIS IS WHAT YOU NEED: ", next(model.parameters()).is_cuda)
    out, train_out = model(img, augment=augment)  # inference and training outputs
    t1 += time_synchronized() - t

    # Compute loss
    if compute_loss:
        loss += compute_loss([x.float() for x in train_out], targets)[1][:3]  # box, obj, cls

    # Run Polygon NMS
    targets[:, 2:] *= torch.Tensor([width, height, width, height, width, height, width, height]).to(device)  # to pixels
    lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else []  # for autolabelling
    t = time_synchronized()
    # out is list of detections, on (n,10) tensor per image [xyxyxyxy, conf, cls]

    # polygon_non_max_suppression takes most of the time to operate
    # For UCAS-AOD dataset, multi_label=False
    # Polygon does not support agnostic
    out = polygon_non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=False)    
    t2 += time_synchronized() - t

    # Statistics per image
    for si, pred in enumerate(out):
        labels = targets[targets[:, 0] == si, 1:]
        nl = len(labels)
        tcls = labels[:, 0].tolist() if nl else []  # target class
        path = Path(paths[si])
        seen += 1

        if len(pred) == 0:
            if nl:
                stats.append((torch.zeros(0, niou, dtype=torch.bool).to('cpu'), torch.Tensor().to('cpu'), torch.Tensor().to('cpu'), tcls))
            continue

        # Predictions
        if single_cls:
            pred[:, 9] = 0
        predn = pred.clone()
        polygon_scale_coords(img[si].shape[1:], predn[:, :8], shapes[si][0], shapes[si][1])  # native-space pred

        # Append to text file
        if save_txt:
            gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0, 1, 0, 1, 0]]  # normalization gain xyxyxyxy
            for *xyxyxyxy, conf, cls in predn.tolist():
                xyxyxyxyn = ((torch.tensor(xyxyxyxy).view(1, 8)) / gn).view(-1).tolist()  # normalized xyxyxyxy
                line = (cls, *xyxyxyxyn, conf) if save_conf else (cls, *xyxyxyxyn)  # label format
                with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
                    f.write(('%g ' * len(line)).rstrip() % line + '\n')

        # W&B logging - Media Panel plots
        if wandb_logger and wandb_logger.wandb:
            # check for test operations
            if ((wandb_logger.current_epoch > 0) and (wandb_logger.bbox_interval > 1) 
                and (wandb_logger.current_epoch % wandb_logger.bbox_interval == 0)):
                if batch_i == 0:
                    if wandb_logger.current_epoch == wandb_logger.bbox_interval:
                        f = save_dir / f'test_batch{batch_i}_labels.jpg'  # predictions
                        polygon_plot_images(img, targets, paths, f, names)
                        current_batches = [wandb_logger.wandb.Image(str(f), caption=f.name)]
                        wandb_logger.log({f"Validation-targeted-labels": current_batches})
                    f = save_dir / f'test_batch{batch_i}_pred_epoch{wandb_logger.current_epoch}.jpg'  # predictions
                    polygon_plot_images(img, polygon_output_to_target(out), paths, f, names)
                    current_batches = [wandb_logger.wandb.Image(str(f), caption=f.name)]
                    wandb_logger.log({f"Validation-epoch{wandb_logger.current_epoch}": current_batches})

        wandb_logger.log_training_progress(predn, path, names) if wandb_logger and wandb_logger.wandb_run else None

        # Append to pycocotools JSON dictionary
        if save_json:
            # [{"image_id": 42, "category_id": 18, "polygon_box": [x1, y1, x2, y2, x3, y3, x4, y4], "score": 0.236}, ...
            image_id = int(path.stem) if path.stem.isnumeric() else path.stem
            box = predn[:, :8]  # xyxyxyxy
            for p, b in zip(pred.tolist(), box.tolist()):
                jdict.append({'image_id': image_id,
                              'category_id': coco91class[int(p[9])] if is_coco else int(p[9]),
                              'polygon_box': [round(x, 3) for x in b],
                              'score': round(p[8], 5)})

        # Assign all predictions as incorrect
        correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device)
        if nl:
            detected = []  # target indices
            tcls_tensor = labels[:, 0]

            # target boxes
            tbox = labels[:, 1:9]
            polygon_scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1])  # native-space labels
            if plots:
                confusion_matrix.process_batch(predn, torch.cat((labels[:, 0:1], tbox), 1))

            # Per target class
            for cls in torch.unique(tcls_tensor):
                ti = (cls == tcls_tensor).nonzero(as_tuple=False).view(-1)  # target indices
                pi = (cls == pred[:, 9]).nonzero(as_tuple=False).view(-1)  # prediction indices

                # Search for detections
                if pi.shape[0]:
                    # Prediction to target ious
                    ious, i = polygon_box_iou(predn[pi, :8], tbox[ti], device=device).max(1)  # best ious, indices

                    # Append detections
                    detected_set = set()
                    for j in (ious.to(device) > iouv[index_ap50].to(device)).nonzero(as_tuple=False).to(device):
                        d = ti[i[j]]  # detected target
                        if d.item() not in detected_set:
                            detected_set.add(d.item())
                            detected.append(d)
                            correct[pi[j]] = ious[j].to(device) > iouv.to(device)  # iou_thres is 1xn
                            if len(detected) == nl:  # all targets already located in image
                                break

        # Append statistics (correct, conf, pcls, tcls)
        stats.append((correct.cpu(), pred[:, 8].cpu(), pred[:, 9].cpu(), tcls))

    # Plot images
    if plots and batch_i < 3:
        f = save_dir / f'test_batch{batch_i}_labels.jpg'  # labels
        Thread(target=polygon_plot_images, args=(img, targets, paths, f, names), daemon=True).start()
        f = save_dir / f'test_batch{batch_i}_pred.jpg'  # predictions
        Thread(target=polygon_plot_images, args=(img, polygon_output_to_target(out), paths, f, names), daemon=True).start()

# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)]  # to numpy
if len(stats) and stats[0].any():
    p, r, ap, f1, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names)
    ap50, ap = ap[:, index_ap50], ap.mean(1)  # AP@0.5, AP@0.5:0.95
    mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()
    nt = np.bincount(stats[3].astype(np.int64), minlength=nc)  # number of targets per class
else:
    nt = torch.zeros(1)

# Print results
pf = '%20s' + '%11i' * 2 + '%11.3g' * 4  # print format
print(pf % ('all', seen, nt.sum(), mp, mr, map50, map))

# Print results per class
if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):
    for i, c in enumerate(ap_class):
        print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))

# Print speeds
t = tuple(x / seen * 1E3 for x in (t0, t1, t2))  # speeds per image
if not training:
    shape = (batch_size, 3, imgsz, imgsz)
    print(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t)

# Plots
if plots:
    confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))
    if wandb_logger and wandb_logger.wandb:
        val_batches = [wandb_logger.wandb.Image(str(f), caption=f.name) for f in sorted(save_dir.glob('test*.jpg'))]
        wandb_logger.log({"Validation": val_batches})

# Save JSON
if save_json and len(jdict):
    w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else ''  # weights
    pred_json = str(save_dir / f"{w}_predictions.json")  # predictions json
    print('\nEvaluating pycocotools mAP... saving %s...' % pred_json)
    with open(pred_json, 'w') as f:
        json.dump(jdict, f)

# Return results
model.float()  # for training
if not training:
    s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
    print(f"Results saved to {save_dir}{s}")
maps = np.zeros(nc) + map
for i, c in enumerate(ap_class):
    maps[c] = ap[i]
return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t

if name == 'main': parser = argparse.ArgumentParser(prog='test.py') parser.add_argument('--data', type=str, default='data/coco128.yaml', help='dataset.yaml path') parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)') parser.add_argument('--batch-size', type=int, default=32, help='batch size') parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold') parser.add_argument('--task', default='val', help='train, val, test, speed or study') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') parser.add_argument('--augment', action='store_true', help='augmented inference') parser.add_argument('--verbose', action='store_true', help='report mAP by class') parser.add_argument('--save-txt', action='store_true', help='save results to .txt') parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to .txt') parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file') parser.add_argument('--project', default='runs/test', help='save to project/name') parser.add_argument('--name', default='exp', help='save to project/name') parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')

opt = parser.parse_args()
assert "polygon" in opt.data.lower(), "polygon_test.py is designed for polygon cases"
opt.save_json |= opt.data.endswith('coco.yaml')
opt.save_txt |= opt.save_hybrid
opt.data = check_file(opt.data)  # check file
print(opt)
check_requirements(exclude=('tensorboard', 'thop'))

if opt.task in ('train', 'val', 'test'):  # run normally
    test(**vars(opt))

elif opt.task == 'speed':  # speed benchmarks
    for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
        test(opt.data, weights=w, batch_size=opt.batch_size, imgsz=opt.imgsz, conf_thres=.25, iou_thres=.45,
             save_json=False, plots=False)

elif opt.task == 'study':  # run over a range of settings and save/plot
    # python test.py --task study --data coco.yaml --iou 0.7 --weights yolov5s.pt yolov5m.pt yolov5l.pt yolov5x.pt
    x = list(range(256, 1536 + 128, 128))  # x axis (image sizes)
    for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]:
        f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt'  # filename to save to
        y = []  # y axis
        for i in x:  # img-size
            print(f'\nRunning {f} point {i}...')
            r, _, t = test(opt.data, weights=w, batch_size=opt.batch_size, imgsz=i, conf_thres=opt.conf_thres,
                           iou_thres=opt.iou_thres, save_json=opt.save_json, plots=False)
            y.append(r + t)  # results and times
        np.savetxt(f, y, fmt='%10.4g')  # save
    os.system('zip -r study.zip study_*.txt')
    plot_study_txt(x=x)  # plot

`