def __init__(self, metrics, classes, config=None): self.config = config self.classes = classes self.metrics = [] for classe in classes: self.metrics.append({metric: [] for metric in metrics}) self.modules = {metric: load_module("neat_eo.metrics." + metric) for metric in metrics} self.n = 0
def main(args): config = load_config(args.config) check_channels(config) check_classes(config) assert torch.cuda.is_available( ), "No GPU support found. Check CUDA and NVidia Driver install." assert torch.distributed.is_nccl_available( ), "No NCCL support found. Check your PyTorch install." world_size = torch.cuda.device_count() args.bs = args.bs if args.bs is not None else math.floor(os.cpu_count() / world_size) args.workers = args.workers if args.workers is not None else args.bs palette, transparency = make_palette( [classe["color"] for classe in config["classes"]]) args.cover = [ tile for tile in tiles_from_csv(os.path.expanduser(args.cover)) ] if args.cover else None args.out = os.path.expanduser(args.out) log = Logs(os.path.join(args.out, "log")) chkpt = torch.load(args.checkpoint, map_location=torch.device("cpu")) chkpt["loader"] = "SemSeg" log.log("neo predict on {} GPUs, with {} workers/GPU and {} tiles/batch". format(world_size, args.workers, args.bs)) log.log("Model {} - UUID: {}".format(chkpt["nn"], chkpt["uuid"])) log.log("---") loader = load_module("neat_eo.loaders.{}".format(chkpt["loader"].lower())) lock_file = os.path.abspath(os.path.join(args.out, str(uuid.uuid1()))) dataset = getattr(loader, chkpt["loader"])( config, chkpt["shape_in"][1:3], args.dataset, args.cover, mode="predict", metatiles=args.metatiles, keep_borders=args.keep_borders, ) mp.spawn(gpu_worker, nprocs=world_size, args=(world_size, lock_file, args, config, dataset, palette, transparency)) if os.path.exists(lock_file): os.remove(lock_file) if not args.no_web_ui and dataset.cover: template = "leaflet.html" if not args.web_ui_template else args.web_ui_template base_url = args.web_ui_base_url if args.web_ui_base_url else "." web_ui(args.out, base_url, dataset.cover, dataset.cover, "png", template)
def gpu_worker(rank, world_size, lock_file, args, config, dataset, palette, transparency): dist.init_process_group(backend="nccl", init_method="file://" + lock_file, world_size=world_size, rank=rank) torch.cuda.set_device(rank) chkpt = torch.load(args.checkpoint, map_location=torch.device(rank)) nn_module = load_module("neat_eo.nn.{}".format(chkpt["nn"].lower())) nn = getattr(nn_module, chkpt["nn"])(chkpt["shape_in"], chkpt["shape_out"], chkpt["encoder"].lower()).to(rank) nn = DistributedDataParallel(nn, device_ids=[rank], find_unused_parameters=True) chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location="cuda:{}".format(rank)) assert nn.module.version == chkpt["model_version"], "Model Version mismatch" nn.load_state_dict(chkpt["state_dict"]) sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank) loader = DataLoader(dataset, batch_size=args.bs, shuffle=False, num_workers=args.workers, sampler=sampler) assert len(loader), "Empty predict dataset directory. Check your path." C, W, H = chkpt["shape_out"] nn.eval() with torch.no_grad(): dataloader = tqdm(loader, desc="Predict", unit="Batch/GPU", ascii=True) if rank == 0 else loader for images, tiles in dataloader: if args.metatiles: N = images.shape[0] qs = int(W / 4) hs = int(W / 2) ts = int(W) # fmt:off probs = np.zeros((N, C, W, H), dtype=np.float) probs[:, :, 0:hs, 0:hs] = nn(images[:, :, 0:ts, 0:ts].to(rank)).data.cpu().numpy()[:, :, qs:-qs, qs:-qs] probs[:, :, 0:hs, hs:] = nn(images[:, :, 0:ts, hs:].to(rank)).data.cpu().numpy()[:, :, qs:-qs, qs:-qs] probs[:, :, hs:, 0:hs] = nn(images[:, :, hs:, 0:ts].to(rank)).data.cpu().numpy()[:, :, qs:-qs, qs:-qs] probs[:, :, hs:, hs:] = nn(images[:, :, hs:, hs:].to(rank)).data.cpu().numpy()[:, :, qs:-qs, qs:-qs] # fmt:on else: probs = nn(images.to(rank)).data.cpu().numpy() for tile, prob in zip(tiles, probs): x, y, z = list(map(int, tile)) mask = np.zeros((W, H), dtype=np.uint8) for c in range(C): mask += np.around(prob[c, :, :]).astype(np.uint8) * c tile_label_to_file(args.out, mercantile.Tile(x, y, z), palette, transparency, mask)
def main(args): try: module = load_module("neat_eo.osm.{}".format(args.type.lower())) except: sys.exit("ERROR: Unavailable --type {}".format(args.type)) if os.path.dirname(os.path.expanduser(args.out)): os.makedirs(os.path.dirname(os.path.expanduser(args.out)), exist_ok=True) osmium_handler = getattr(module, "{}Handler".format(args.type))() print("neo extract {} from {} to {}".format(args.type, args.pbf, args.out), file=sys.stderr, flush=True) print("\nNOTICE: could take time. Be patient...\n".format( args.type, args.pbf, args.out), file=sys.stderr, flush=True) osmium_handler.apply_file(filename=os.path.expanduser(args.pbf), locations=True) osmium_handler.save(os.path.expanduser(args.out))
def main(args): config = load_config(args.config) args.out = os.path.expanduser(args.out) args.cover = [ tile for tile in tiles_from_csv(os.path.expanduser(args.cover)) ] if args.cover else None if args.classes_weights: try: args.classes_weights = list( map(float, args.classes_weights.split(","))) except: assert args.classes_weights == "auto", "invalid --classes_weights value" else: args.classes_weights = [ classe["weight"] for classe in config["classes"] ] args.tiles_weights = ([(tile, weight) for tile, weight in tiles_from_csv( os.path.expanduser(args.tiles_weights), extra_columns=True)] if args.tiles_weights else None) config["model"][ "loader"] = args.loader if args.loader else config["model"]["loader"] config["model"]["ts"] = tuple(map( int, args.ts.split(","))) if args.ts else config["model"]["ts"] config["model"]["nn"] = args.nn if args.nn else config["model"]["nn"] config["model"]["encoder"] = args.encoder if args.encoder else config[ "model"]["encoder"] config["train"]["bs"] = args.bs if args.bs else config["train"]["bs"] config["train"][ "loss"] = args.loss if args.loss else config["train"]["loss"] config["train"]["optimizer"][ "name"] = args.optimizer if args.optimizer else config["train"][ "optimizer"]["name"] config["train"]["optimizer"][ "lr"] = args.lr if args.lr else config["train"]["optimizer"]["lr"] check_classes(config) check_channels(config) check_model(config) log = Logs(os.path.join(args.out, "log")) assert torch.cuda.is_available( ), "No GPU support found. Check CUDA and NVidia Driver install." assert torch.distributed.is_nccl_available( ), "No NCCL support found. Check your PyTorch install." world_size = torch.cuda.device_count() if args.train_dataset else 1 args.workers = min( config["train"]["bs"] if not args.workers else args.workers, math.floor(os.cpu_count() / world_size)) assert args.eval_dataset or args.train_dataset, "Provide at least one dataset" if args.eval_dataset and not args.train_dataset and not args.checkpoint: log.log( "\n\nNOTICE: No Checkpoint provided for eval only. Seems peculiar.\n\n" ) log.log("neo train/eval on {} GPUs, with {} workers/GPU".format( world_size, args.workers)) log.log("---") loader = load_module("neat_eo.loaders.{}".format( config["model"]["loader"].lower())) train_dataset = None if args.train_dataset: assert os.path.isdir(os.path.expanduser( args.train_dataset)), "--train_dataset path is not a directory" train_dataset = getattr(loader, config["model"]["loader"])( config, config["model"]["ts"], args.train_dataset, args.cover, args.tiles_weights, "train") assert len(train_dataset), "Empty or Invalid --train_dataset content" shape_in = train_dataset.shape_in shape_out = train_dataset.shape_out log.log("\nDataSet Training: {}".format(args.train_dataset)) if args.classes_weights == "auto": args.classes_weights = compute_classes_weights( args.train_dataset, config["classes"], args.cover, os.cpu_count()) eval_dataset = None if args.eval_dataset: assert os.path.isdir(os.path.expanduser( args.eval_dataset)), "--eval_dataset path is not a directory" eval_dataset = getattr(loader, config["model"]["loader"])( config, config["model"]["ts"], args.eval_dataset, args.cover, args.tiles_weights, "eval") assert len(eval_dataset), "Empty or Invalid --eval_dataset content" shape_in = eval_dataset.shape_in shape_out = eval_dataset.shape_out log.log("DataSet Eval: {}".format(args.eval_dataset)) if not args.train_dataset and args.classes_weights == "auto": args.classes_weights = compute_classes_weights( args.eval_dataset, config["classes"], args.cover, os.cpu_count()) log.log("\n--- Input tensor") num_channel = 1 # 1-based numerotation for channel in config["channels"]: for band in channel["bands"]: log.log("Channel {}:\t\t {} - (band:{})".format( num_channel, channel["name"], band)) num_channel += 1 log.log("\n--- Output Classes ---") for c, classe in enumerate(config["classes"]): log.log("Class {}:\t\t {} ({:.2f})".format(c, classe["title"], args.classes_weights[c])) log.log("\n--- Model ---") for hp in config["model"]: log.log("{}{}".format(hp.ljust(25, " "), config["model"][hp])) lock_file = os.path.abspath(os.path.join(args.out, str(uuid.uuid1()))) mp.spawn( gpu_worker, nprocs=world_size, args=(world_size, lock_file, train_dataset, eval_dataset, shape_in, shape_out, args, config), ) if os.path.exists(lock_file): os.remove(lock_file)
def gpu_worker(rank, world_size, lock_file, train_dataset, eval_dataset, shape_in, shape_out, args, config): log = Logs(os.path.join(args.out, "log")) if rank == 0 else None csv_train = open(os.path.join(args.out, "train.csv"), mode="a") if train_dataset and rank == 0 else None csv_eval = open(os.path.join(args.out, "eval.csv"), mode="a") if eval_dataset and rank == 0 else None dist.init_process_group(backend="nccl", init_method="file://" + lock_file, world_size=world_size, rank=rank) torch.cuda.set_device(rank) torch.manual_seed(0) bs = config["train"]["bs"] if train_dataset: sampler = torch.utils.data.distributed.DistributedSampler( train_dataset, num_replicas=world_size, rank=rank) train_loader = DataLoader(train_dataset, batch_size=bs, shuffle=False, drop_last=True, num_workers=args.workers, sampler=sampler) else: train_loader = None if eval_dataset: eval_loader = DataLoader(eval_dataset, batch_size=bs, shuffle=False, drop_last=True, num_workers=args.workers) else: eval_loader = None nn_module = load_module("neat_eo.nn.{}".format( config["model"]["nn"].lower())) nn = getattr(nn_module, config["model"]["nn"])(shape_in, shape_out, config["model"]["encoder"].lower(), config["train"]).cuda(rank) nn = DistributedDataParallel(nn, device_ids=[rank], find_unused_parameters=True) if train_dataset: optimizer_params = { key: value for key, value in config["train"]["optimizer"].items() if key != "name" } optimizer = getattr(torch.optim, config["train"]["optimizer"]["name"])( nn.parameters(), **optimizer_params) if rank == 0: log.log("\n--- Train ---") for hp in config["train"]: if hp == "da": da = config["train"]["da"]["name"] dap = config["train"]["da"]["p"] log.log("{}{} ({:.2f})".format("da".ljust(25, " "), da, dap)) elif hp == "metrics": log.log("{}{}".format(hp.ljust( 25, " "), set(config["train"][hp]))) # aesthetic elif hp != "optimizer": log.log("{}{}".format(hp.ljust(25, " "), config["train"][hp])) log.log("{}{}".format("optimizer".ljust(25, " "), config["train"]["optimizer"]["name"])) for k, v in optimizer.state_dict()["param_groups"][0].items(): if k != "params": log.log(" - {}{}".format(k.ljust(25 - 3, " "), v)) resume = 0 if args.checkpoint: chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location="cuda:{}".format(rank)) assert nn.module.version == chkpt[ "model_version"], "Model Version mismatch" nn.load_state_dict(chkpt["state_dict"]) if rank == 0: log.log("\n--- Using Checkpoint ---") log.log("Path:\t\t {}".format(args.checkpoint)) log.log("UUID:\t\t {}".format(chkpt["uuid"])) if args.resume: optimizer.load_state_dict(chkpt["optimizer"]) resume = chkpt["epoch"] assert resume < args.epochs, "Epoch asked, already reached by the given checkpoint" loss_module = load_module("neat_eo.losses.{}".format( config["train"]["loss"].lower())) criterion = getattr(loss_module, config["train"]["loss"])().cuda(rank) if eval_dataset and not train_dataset: do_epoch(rank, eval_loader, config, args.classes_weights, log, csv_eval, nn, criterion, "eval", 1) dist.destroy_process_group() return for epoch in range(resume + 1, args.epochs + 1): # 1-N based if train_dataset: if rank == 0: log.log("\n---\nEpoch: {}/{}\n".format(epoch, args.epochs)) sampler.set_epoch( epoch) # https://github.com/pytorch/pytorch/issues/31232 do_epoch(rank, train_loader, config, args.classes_weights, log, csv_train, nn, criterion, "train", epoch, optimizer) if rank == 0: UUID = uuid.uuid1() states = { "uuid": UUID, "model_version": nn.module.version, "producer_name": "Neat-EO.pink", "producer_version": neo.__version__, "model_licence": "MIT", "domain": "pink.Neat-EO", # reverse-DNS "doc_string": nn.module.doc_string, "shape_in": shape_in, "shape_out": shape_out, "state_dict": nn.state_dict(), "epoch": epoch, "nn": config["model"]["nn"], "encoder": config["model"]["encoder"], "optimizer": optimizer.state_dict(), "loader": config["model"]["loader"], } checkpoint_path = os.path.join( args.out, "checkpoint-{:05d}.pth".format(epoch)) if epoch == args.epochs or not (epoch % args.saving): log.log("\n--- Saving Checkpoint ---") log.log("Path:\t\t {}".format(checkpoint_path)) log.log("UUID:\t\t {}\n".format(UUID)) torch.save(states, checkpoint_path) dist.barrier() if eval_dataset: do_epoch(rank, eval_loader, config, args.classes_weights, log, csv_eval, nn, criterion, "eval", epoch) dist.destroy_process_group()
def main(args): chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location=torch.device("cpu")) assert chkpt, "Unable to load checkpoint {}".format(args.checkpoint) if os.path.dirname(os.path.expanduser(args.out)): os.makedirs(os.path.dirname(os.path.expanduser(args.out)), exist_ok=True) args.out = os.path.expanduser(args.out) UUID = chkpt["uuid"] if "uuid" in chkpt else uuid.uuid1() try: nn_name = chkpt["nn"] except: assert args.nn, "--nn mandatory as not already in input .pth" nn_name = args.nn try: encoder = chkpt["encoder"] except: assert args.encoder, "--encoder mandatory as not already in input .pth" encoder = args.encoder try: loader = chkpt["loader"] except: assert args.loader, "--loader mandatory as not already in input .pth" doc_string = args.doc_string try: doc_string = chkpt["doc_string"] except: assert args.doc_string, "--doc_string mandatory as not already in input .pth" doc_string = args.doc_string try: shape_in = chkpt["shape_in"] except: assert args.shape_in, "--shape_in mandatory as not already in input .pth" shape_in = tuple(map(int, args.shape_in.split(","))) try: shape_out = chkpt["shape_out"] except: assert args.shape_out, "--shape_out mandatory as not already in input .pth" shape_out = tuple(map(int, args.shape_out.split(","))) nn_module = load_module("neat_eo.nn.{}".format(nn_name.lower())) nn = getattr(nn_module, nn_name)(shape_in, shape_out, encoder.lower()).to("cpu") print("neo export model to {}".format(args.type), file=sys.stderr) print("Model: {}".format(nn_name, file=sys.stderr)) print("UUID: {}".format(UUID, file=sys.stderr)) if args.type == "pth": states = { "uuid": UUID, "model_version": None, "producer_name": "Neat-EO.pink", "producer_version": neo.__version__, "model_licence": "MIT", "domain": "pink.Neat-EO", # reverse-DNS "doc_string": doc_string, "shape_in": shape_in, "shape_out": shape_out, "state_dict": nn.state_dict(), "epoch": 0, "nn": nn_name, "encoder": encoder, "optimizer": None, "loader": loader, } torch.save(states, args.out) else: try: # https://github.com/pytorch/pytorch/issues/9176 nn.module.state_dict(chkpt["state_dict"]) except AttributeError: nn.state_dict(chkpt["state_dict"]) nn.eval() batch = torch.rand(1, *shape_in) if args.type == "onnx": torch.onnx.export( nn, torch.autograd.Variable(batch), args.out, input_names=["input", "shape_in", "shape_out"], output_names=["output"], dynamic_axes={ "input": { 0: "num_batch" }, "output": { 0: "num_batch" } }, ) if args.type == "jit": torch.jit.trace(nn, batch).save(args.out)
def main(args): if not args.masks or not args.labels: assert args.mode != "list", "Parameters masks and labels are mandatories in list mode." assert not (args.min or args.max), "Both --masks and --labels mandatory, for metric filtering." if args.min or args.max: config = load_config(args.config) args.out = os.path.expanduser(args.out) cover = [tile for tile in tiles_from_csv(os.path.expanduser(args.cover))] if args.cover else None args_minmax = set() args.min = {(m[0], m[1]): m[2] for m in args.min} if args.min else dict() args.max = {(m[0], m[1]): m[2] for m in args.max} if args.max else dict() args_minmax.update(args.min.keys()) args_minmax.update(args.max.keys()) minmax = dict() for mm in args_minmax: mm_min = float(args.min[mm]) if mm in args.min else 0.0 mm_max = float(args.max[mm]) if mm in args.max else 1.0 assert mm_min < mm_max, "--min must be lower than --max, on {}".format(mm) minmax[mm] = { "min": mm_min, "max": mm_max, "class_id": [c for c, classe in enumerate(config["classes"]) if classe["title"] == mm[0]][0], "module": load_module("neat_eo.metrics." + mm[1]), } if not args.workers: args.workers = os.cpu_count() print("neo compare {} on CPU, with {} workers".format(args.mode, args.workers), file=sys.stderr, flush=True) if args.images: tiles = [tile for tile in tiles_from_dir(args.images[0], cover=cover)] assert len(tiles), "Empty images dir: {}".format(args.images[0]) for image in args.images[1:]: assert sorted(tiles) == sorted([tile for tile in tiles_from_dir(image, cover=cover)]), "Unconsistent images dirs" if args.labels and args.masks: tiles_masks = [tile for tile in tiles_from_dir(args.masks, cover=cover)] tiles_labels = [tile for tile in tiles_from_dir(args.labels, cover=cover)] if args.images: assert sorted(tiles) == sorted(tiles_masks) == sorted(tiles_labels), "Unconsistent images/label/mask directories" else: assert len(tiles_masks), "Empty masks dir: {}".format(args.masks) assert len(tiles_labels), "Empty labels dir: {}".format(args.labels) assert sorted(tiles_masks) == sorted(tiles_labels), "Label and Mask directories are not consistent" tiles = tiles_masks tiles_list = [] tiles_compare = [] progress = tqdm(total=len(tiles), ascii=True, unit="tile") log = False if args.mode == "list" else Logs(os.path.join(args.out, "log")) with futures.ThreadPoolExecutor(args.workers) as executor: def worker(tile): x, y, z = list(map(str, tile)) if args.masks and args.labels: label = np.array(Image.open(os.path.join(args.labels, z, x, "{}.png".format(y)))) mask = np.array(Image.open(os.path.join(args.masks, z, x, "{}.png".format(y)))) assert label.shape == mask.shape, "Inconsistent tiles (size or dimensions)" metrics = dict() for mm in minmax: try: metrics[mm] = getattr(minmax[mm]["module"], "get")( torch.as_tensor(label, device="cpu"), torch.as_tensor(mask, device="cpu"), minmax[mm]["class_id"], ) except: progress.update() return False, tile if not (minmax[mm]["min"] <= metrics[mm] <= minmax[mm]["max"]): progress.update() return True, tile tiles_compare.append(tile) if args.mode == "side": for i, root in enumerate(args.images): img = tile_image_from_file(tile_from_xyz(root, x, y, z)[1], force_rgb=True) if i == 0: side = np.zeros((img.shape[0], img.shape[1] * len(args.images), 3)) side = np.swapaxes(side, 0, 1) if args.vertical else side image_shape = img.shape else: assert image_shape[0:2] == img.shape[0:2], "Unconsistent image size to compare" if args.vertical: side[i * image_shape[0] : (i + 1) * image_shape[0], :, :] = img else: side[:, i * image_shape[0] : (i + 1) * image_shape[0], :] = img tile_image_to_file(args.out, tile, np.uint8(side)) elif args.mode == "stack": for i, root in enumerate(args.images): tile_image = tile_image_from_file(tile_from_xyz(root, x, y, z)[1], force_rgb=True) if i == 0: image_shape = tile_image.shape[0:2] stack = tile_image / len(args.images) else: assert image_shape == tile_image.shape[0:2], "Unconsistent image size to compare" stack = stack + (tile_image / len(args.images)) tile_image_to_file(args.out, tile, np.uint8(stack)) elif args.mode == "list": tiles_list.append([tile, metrics]) progress.update() return True, tile for ok, tile in executor.map(worker, tiles): if not ok and log: log.log("Warning: skipping. {}".format(str(tile))) if args.mode == "list": with open(args.out, mode="w") as out: if args.geojson: out.write('{"type":"FeatureCollection","features":[') first = True for tile_list in tiles_list: tile, metrics = tile_list x, y, z = list(map(str, tile)) if args.geojson: prop = '"properties":{{"x":{},"y":{},"z":{}'.format(x, y, z) for metric in metrics: prop += ',"{}":{:.3f}'.format(metric, metrics[metric]) geom = '"geometry":{}'.format(json.dumps(feature(tile, precision=6)["geometry"])) out.write('{}{{"type":"Feature",{},{}}}}}'.format("," if not first else "", geom, prop)) first = False if not args.geojson: out.write("{},{},{}".format(x, y, z)) for metric in metrics: out.write("\t{:.3f}".format(metrics[metric])) out.write(os.linesep) if args.geojson: out.write("]}") out.close() base_url = args.web_ui_base_url if args.web_ui_base_url else "." if args.mode == "side" and not args.no_web_ui: template = "compare.html" if not args.web_ui_template else args.web_ui_template web_ui(args.out, base_url, tiles, tiles_compare, args.format, template, union_tiles=False) if args.mode == "stack" and not args.no_web_ui: template = "leaflet.html" if not args.web_ui_template else args.web_ui_template tiles = [tile for tile in tiles_from_dir(args.images[0])] web_ui(args.out, base_url, tiles, tiles_compare, args.format, template)
def __init__(self, shape_in, shape_out, encoder="resnet50", train_config=None): super().__init__() doc = "U-Net like encoder-decoder architecture with a ResNet, ResNext or WideResNet encoder.\n\n" doc += " - https://arxiv.org/abs/1505.04597 - U-Net: Convolutional Networks for Biomedical Image Segmentation\n" if encoder in ["resnet50", "resnet101", "resnet152"]: doc += " - https://arxiv.org/abs/1512.03385 - Deep Residual Learning for Image Recognition\n" elif encoder in ["resnext50_32x4d", "resnext101_32x8d"]: doc += " - https://arxiv.org/pdf/1611.05431 - Aggregated Residual Transformations for DNN\n" elif encoder in ["wide_resnet50_2", "wide_resnet101_2"]: doc += " - https://arxiv.org/abs/1605.07146 - Wide Residual Networks\n" else: encoders = "resnet50, resnet101, resnet152, resnext50_32x4d, resnext101_32x8d, wide_resnet50_2, wide_resnet101_2" assert False, "Albunet, expects as encoder: " + encoders self.version = 2 self.doc_string = doc num_filters = 32 num_channels = shape_in[0] num_classes = shape_out[0] assert num_channels, "Empty Channels" assert num_classes, "Empty Classes" try: pretrained = train_config["pretrained"] except: pretrained = False models = load_module("torchvision.models") self.encoder = getattr(models, encoder)(pretrained=pretrained) # https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py if num_channels != 3: weights = nn.init.kaiming_normal_(torch.zeros( (64, num_channels, 7, 7)), mode="fan_out", nonlinearity="relu") if pretrained: for c in range(min(num_channels, 3)): weights.data[:, c, :, :] = self.encoder.conv1.weight.data[:, c, :, :] self.encoder.conv1 = nn.Conv2d(num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False) self.encoder.conv1.weight = nn.Parameter(weights) self.center = DecoderBlock(2048, num_filters * 8) self.dec0 = DecoderBlock(2048 + num_filters * 8, num_filters * 8) self.dec1 = DecoderBlock(1024 + num_filters * 8, num_filters * 8) self.dec2 = DecoderBlock(512 + num_filters * 8, num_filters * 2) self.dec3 = DecoderBlock(256 + num_filters * 2, num_filters * 2 * 2) self.dec4 = DecoderBlock(num_filters * 2 * 2, num_filters) self.dec5 = ConvRelu(num_filters, num_filters) self.final = nn.Conv2d(num_filters, num_classes, kernel_size=1)
def main(args): config = load_config(args.config) args.cover = [ tile for tile in tiles_from_csv(os.path.expanduser(args.cover)) ] if args.cover else None if args.classes_weights: try: args.classes_weights = list( map(float, args.classes_weights.split(","))) except: assert args.classes_weights == "auto", "invalid --classes_weights value" args.classes_weights = compute_classes_weights( args.dataset, config["classes"], args.cover, os.cpu_count()) else: args.classes_weights = [ classe["weight"] for classe in config["classes"] ] args.tiles_weights = ([(tile, weight) for tile, weight in tiles_from_csv( os.path.expanduser(args.tiles_weights), extra_columns=True)] if args.tiles_weights else None) args.bs = args.bs if args.bs else config["train"]["bs"] check_classes(config) check_channels(config) check_model(config) assert torch.cuda.is_available( ), "No GPU support found. Check CUDA and NVidia Driver install." assert torch.distributed.is_nccl_available( ), "No NCCL support found. Check your PyTorch install." world_size = 1 # Hard Coded since eval MultiGPUs not yet implemented args.workers = min(args.bs if not args.workers else args.workers, math.floor(os.cpu_count() / world_size)) print("neo eval on 1 GPU, with {} workers, and {} tiles/batch".format( args.workers, args.bs)) loader = load_module("neat_eo.loaders.{}".format( config["model"]["loader"].lower())) assert os.path.isdir(os.path.expanduser( args.dataset)), "--dataset path is not a directory" dataset = getattr(loader, config["model"]["loader"])(config, config["model"]["ts"], args.dataset, args.cover, args.tiles_weights, "eval") assert len(dataset), "Empty or Invalid --dataset content" shape_in = dataset.shape_in shape_out = dataset.shape_out print("DataSet Eval: {}".format(args.dataset)) print("\n--- Input tensor") num_channel = 1 # 1-based numerotation for channel in config["channels"]: for band in channel["bands"]: print("Channel {}:\t\t {} - (band:{})".format( num_channel, channel["name"], band)) num_channel += 1 print("\n--- Output Classes ---") for c, classe in enumerate(config["classes"]): print("Class {}:\t\t {} ({:.2f})".format(c, classe["title"], args.classes_weights[c])) print("\n--- Model ---") for hp in config["model"]: print("{}{}".format(hp.ljust(25, " "), config["model"][hp])) lock_file = os.path.abspath(os.path.join("/tmp", str(uuid.uuid1()))) mp.spawn(gpu_worker, nprocs=world_size, args=(world_size, lock_file, dataset, shape_in, shape_out, args, config)) if os.path.exists(lock_file): os.remove(lock_file)
def gpu_worker(rank, world_size, lock_file, dataset, shape_in, shape_out, args, config): dist.init_process_group(backend="nccl", init_method="file://" + lock_file, world_size=world_size, rank=rank) torch.cuda.set_device(rank) torch.manual_seed(0) loader = DataLoader(dataset, batch_size=args.bs, shuffle=False, drop_last=True, num_workers=args.workers) nn_module = load_module("neat_eo.nn.{}".format( config["model"]["nn"].lower())) nn = getattr(nn_module, config["model"]["nn"])(shape_in, shape_out, config["model"]["encoder"].lower(), config["train"]).cuda(rank) nn = DistributedDataParallel(nn, device_ids=[rank], find_unused_parameters=True) if args.checkpoint: chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location="cuda:{}".format(rank)) assert nn.module.version == chkpt[ "model_version"], "Model Version mismatch" nn.load_state_dict(chkpt["state_dict"]) if rank == 0: print("\n--- Using Checkpoint ---") print("Path:\t\t {}".format(args.checkpoint)) print("UUID:\t\t {}".format(chkpt["uuid"])) nn.eval() with torch.no_grad(): args.metrics = args.metrics if args.metrics else config["train"][ "metrics"] metrics = Metrics(args.metrics, config["classes"], config=config) assert len(loader), "Empty or Inconsistent DataSet" dataloader = tqdm(loader, desc="Eval", unit="Batch", ascii=True) if rank == 0 else loader for images, masks, tiles, tiles_weights in dataloader: images = images.cuda(rank, non_blocking=True) masks = masks.cuda(rank, non_blocking=True) outputs = nn(images) if rank == 0: for mask, output in zip(masks, outputs): metrics.add(mask, output) if rank == 0: print("\n{} μ\t σ".format(" ".ljust(25, " "))) for c, classe in enumerate(config["classes"]): if classe["weight"] != 0.0 and classe["color"] != "transparent": for k, v in metrics.get()[c].items(): print("{}{:.3f}\t {:.3f}".format( (classe["title"] + " " + k).ljust(25, " "), v["μ"], v["σ"])) dist.destroy_process_group()