def main(args):
config = load_config(args.config)
args.out = os.path.expanduser(args.out)
config["model"][
"loader"] = args.loader if args.loader else config["model"]["loader"]
config["model"]["bs"] = args.bs if args.bs else config["model"]["bs"]
config["model"]["lr"] = args.lr if args.lr else config["model"]["lr"]
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["model"][
"loss"] = args.loss if args.loss else config["model"]["loss"]
config["model"]["da"] = args.da if args.da else config["model"]["da"]
config["model"]["dap"] = args.dap if args.dap else config["model"]["dap"]
args.workers = config["model"]["bs"] if not args.workers else args.workers
check_classes(config)
check_channels(config)
check_model(config)
assert os.path.isdir(os.path.expanduser(
args.dataset)), "Dataset is not a directory"
if args.no_training and args.no_validation:
sys.exit()
log = Logs(os.path.join(args.out, "log"))
csv_train = None if args.no_training else open(
os.path.join(args.out, "training.csv"), mode="a")
csv_val = None if args.no_validation else open(
os.path.join(args.out, "validation.csv"), mode="a")
if torch.cuda.is_available():
log.log("RoboSat.pink - training on {} GPUs, with {} workers".format(
torch.cuda.device_count(), args.workers))
log.log("(Torch:{} Cuda:{} CudNN:{})".format(
torch.__version__, torch.version.cuda,
torch.backends.cudnn.version()))
device = torch.device("cuda")
torch.backends.cudnn.benchmark = True
else:
log.log("RoboSat.pink - training on CPU, with {} workers - (Torch:{})".
format(args.workers, torch.__version__))
log.log("")
log.log("==========================================================")
log.log("WARNING: Are you -really- sure about not training on GPU ?")
log.log("==========================================================")
log.log("")
device = torch.device("cpu")
log.log("--- Input tensor from Dataset: {} ---".format(args.dataset))
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("--- Output Classes ---")
for c, classe in enumerate(config["classes"]):
log.log("Class {}:\t\t {}".format(c, classe["title"]))
log.log("--- Hyper Parameters ---")
for hp in config["model"]:
log.log("{}{}".format(hp.ljust(25, " "), config["model"][hp]))
loader = load_module("robosat_pink.loaders.{}".format(
config["model"]["loader"].lower()))
loader_train = getattr(loader, config["model"]["loader"])(
config, config["model"]["ts"], os.path.join(args.dataset,
"training"), None, "train")
loader_val = getattr(loader, config["model"]["loader"])(
config, config["model"]["ts"],
os.path.join(args.dataset, "validation"), None, "train")
encoder = config["model"]["encoder"].lower()
nn_module = load_module("robosat_pink.nn.{}".format(
config["model"]["nn"].lower()))
nn = getattr(nn_module, config["model"]["nn"])(loader_train.shape_in,
loader_train.shape_out,
encoder, config).to(device)
nn = torch.nn.DataParallel(nn)
optimizer = Adam(nn.parameters(), lr=config["model"]["lr"])
resume = 0
if args.checkpoint:
chkpt = torch.load(os.path.expanduser(args.checkpoint),
map_location=device)
nn.load_state_dict(chkpt["state_dict"])
log.log("--- 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("robosat_pink.losses.{}".format(
config["model"]["loss"].lower()))
criterion = getattr(loss_module, config["model"]["loss"])().to(device)
bs = config["model"]["bs"]
train_loader = DataLoader(loader_train,
batch_size=bs,
shuffle=True,
drop_last=True,
num_workers=args.workers)
val_loader = DataLoader(loader_val,
batch_size=bs,
shuffle=False,
drop_last=True,
num_workers=args.workers)
if args.no_training:
epoch = 0
process(val_loader, config, log, csv_val, epoch, device, nn, criterion,
"eval")
sys.exit()
for epoch in range(resume + 1, args.epochs + 1): # 1-N based
UUID = uuid.uuid1()
log.log("---{}Epoch: {}/{} -- UUID: {}".format(os.linesep, epoch,
args.epochs, UUID))
process(train_loader, config, log, csv_train, epoch, device, nn,
criterion, "train", optimizer)
try: # https://github.com/pytorch/pytorch/issues/9176
nn_doc = nn.module.doc
nn_version = nn.module.version
except AttributeError:
nn_doc = nn.doc
nn_version = nn.version
states = {
"uuid": UUID,
"model_version": nn_version,
"producer_name": "RoboSat.pink",
"producer_version": rsp.__version__,
"model_licence": "MIT",
"domain": "pink.RoboSat", # reverse-DNS
"doc_string": nn_doc,
"shape_in": loader_train.shape_in,
"shape_out": loader_train.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("[Saving checkpoint]")
torch.save(states, checkpoint_path)
if not args.no_validation:
process(val_loader, config, log, csv_val, epoch, device, nn,
criterion, "eval")
def main(args):
config = load_config(args.config)
args.out = os.path.expanduser(args.out)
args.workers = torch.cuda.device_count() * 2 if torch.device("cuda") and not args.workers else args.workers
config["model"]["loader"] = args.loader if args.loader else config["model"]["loader"]
config["model"]["bs"] = args.bs if args.bs else config["model"]["bs"]
config["model"]["lr"] = args.lr if args.lr else config["model"]["lr"]
config["model"]["ts"] = args.ts if args.ts else config["model"]["ts"]
config["model"]["nn"] = args.nn if args.nn else config["model"]["nn"]
config["model"]["loss"] = args.loss if args.loss else config["model"]["loss"]
config["model"]["da"] = args.da if args.da else config["model"]["da"]
config["model"]["dap"] = args.dap if args.dap else config["model"]["dap"]
check_classes(config)
check_channels(config)
check_model(config)
if not os.path.isdir(os.path.expanduser(args.dataset)):
sys.exit("ERROR: dataset {} is not a directory".format(args.dataset))
log = Logs(os.path.join(args.out, "log"))
if torch.cuda.is_available():
log.log("RoboSat.pink - training on {} GPUs, with {} workers".format(torch.cuda.device_count(), args.workers))
log.log("(Torch:{} Cuda:{} CudNN:{})".format(torch.__version__, torch.version.cuda, torch.backends.cudnn.version()))
device = torch.device("cuda")
torch.backends.cudnn.benchmark = True
else:
log.log("RoboSat.pink - training on CPU, with {} workers - (Torch:{})".format(args.workers, torch.__version__))
log.log("WARNING: Are you really sure sure about not training on GPU ?")
device = torch.device("cpu")
loader = load_module("robosat_pink.loaders.{}".format(config["model"]["loader"].lower()))
loader_train = getattr(loader, config["model"]["loader"])(
config, config["model"]["ts"], os.path.join(args.dataset, "training"), "train"
)
loader_val = getattr(loader, config["model"]["loader"])(
config, config["model"]["ts"], os.path.join(args.dataset, "validation"), "train"
)
model_module = load_module("robosat_pink.models.{}".format(config["model"]["nn"].lower()))
nn = getattr(model_module, config["model"]["nn"])(loader_train.shape_in, loader_train.shape_out, config).to(device)
nn = torch.nn.DataParallel(nn)
optimizer = Adam(nn.parameters(), lr=config["model"]["lr"])
resume = 0
if args.checkpoint:
chkpt = torch.load(os.path.expanduser(args.checkpoint), map_location=device)
nn.load_state_dict(chkpt["state_dict"])
log.log("Using checkpoint: {}".format(args.checkpoint))
if args.resume:
optimizer.load_state_dict(chkpt["optimizer"])
resume = chkpt["epoch"]
if resume >= args.epochs:
sys.exit("ERROR: Epoch {} already reached by the given checkpoint".format(config["model"]["epochs"]))
loss_module = load_module("robosat_pink.losses.{}".format(config["model"]["loss"].lower()))
criterion = getattr(loss_module, config["model"]["loss"])().to(device)
bs = config["model"]["bs"]
train_loader = DataLoader(loader_train, batch_size=bs, shuffle=True, drop_last=True, num_workers=args.workers)
val_loader = DataLoader(loader_val, batch_size=bs, shuffle=False, drop_last=True, num_workers=args.workers)
log.log("--- Input tensor from Dataset: {} ---".format(args.dataset))
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("--- Hyper Parameters ---")
for hp in config["model"]:
log.log("{}{}".format(hp.ljust(25, " "), config["model"][hp]))
for epoch in range(resume, args.epochs):
UUID = uuid.uuid1()
log.log("---{}Epoch: {}/{} -- UUID: {}".format(os.linesep, epoch + 1, args.epochs, UUID))
process(train_loader, config, log, device, nn, criterion, "train", optimizer)
if not args.no_validation:
process(val_loader, config, log, device, nn, criterion, "eval")
try: # https://github.com/pytorch/pytorch/issues/9176
nn_doc = nn.module.doc
nn_version = nn.module.version
except AttributeError:
nn_version = nn.version
nn_doc == nn.doc
states = {
"uuid": UUID,
"model_version": nn_version,
"producer_name": "RoboSat.pink",
"producer_version": "0.4.0",
"model_licence": "MIT",
"domain": "pink.RoboSat", # reverse-DNS
"doc_string": nn_doc,
"shape_in": loader_train.shape_in,
"shape_out": loader_train.shape_out,
"state_dict": nn.state_dict(),
"epoch": epoch + 1,
"nn": config["model"]["nn"],
"optimizer": optimizer.state_dict(),
"loader": config["model"]["loader"],
}
checkpoint_path = os.path.join(args.out, "checkpoint-{:05d}.pth".format(epoch + 1))
torch.save(states, checkpoint_path)
def main(args):
config = load_config(args.config)
check_channels(config)
check_classes(config)
palette = make_palette([classe["color"] for classe in config["classes"]])
args.workers = torch.cuda.device_count() * 2 if torch.device(
"cuda") and not args.workers else args.workers
cover = [tile for tile in tiles_from_csv(os.path.expanduser(args.cover))
] if args.cover else None
log = Logs(os.path.join(args.out, "log"))
if torch.cuda.is_available():
log.log("RoboSat.pink - predict on {} GPUs, with {} workers".format(
torch.cuda.device_count(), args.workers))
log.log("(Torch:{} Cuda:{} CudNN:{})".format(
torch.__version__, torch.version.cuda,
torch.backends.cudnn.version()))
device = torch.device("cuda")
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
else:
log.log("RoboSat.pink - predict on CPU, with {} workers".format(
args.workers))
log.log("")
log.log("============================================================")
log.log("WARNING: Are you -really- sure about not predicting on GPU ?")
log.log("============================================================")
log.log("")
device = torch.device("cpu")
chkpt = torch.load(args.checkpoint, map_location=device)
model_module = load_module("robosat_pink.models.{}".format(
chkpt["nn"].lower()))
nn = getattr(model_module, chkpt["nn"])(chkpt["shape_in"],
chkpt["shape_out"]).to(device)
nn = torch.nn.DataParallel(nn)
nn.load_state_dict(chkpt["state_dict"])
nn.eval()
log.log("Model {} - UUID: {}".format(chkpt["nn"], chkpt["uuid"]))
mode = "predict" if not args.translate else "predict_translate"
loader_module = load_module("robosat_pink.loaders.{}".format(
chkpt["loader"].lower()))
loader_predict = getattr(loader_module,
chkpt["loader"])(config,
chkpt["shape_in"][1:3],
args.dataset,
cover,
mode=mode)
loader = DataLoader(loader_predict,
batch_size=args.bs,
num_workers=args.workers)
assert len(loader), "Empty predict dataset directory. Check your path."
tiled = []
with torch.no_grad(
): # don't track tensors with autograd during prediction
for images, tiles in tqdm(loader,
desc="Eval",
unit="batch",
ascii=True):
images = images.to(device)
outputs = nn(images)
probs = torch.nn.functional.softmax(outputs,
dim=1).data.cpu().numpy()
for tile, prob in zip(tiles, probs):
x, y, z = list(map(int, tile))
mask = np.around(prob[1:, :, :]).astype(np.uint8).squeeze()
if args.translate:
tile_translate_to_file(args.out, mercantile.Tile(x, y, z),
palette, mask)
else:
tile_label_to_file(args.out, mercantile.Tile(x, y, z),
palette, mask)
tiled.append(mercantile.Tile(x, y, z))
if not args.no_web_ui and not args.translate:
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, tiled, tiled, "png", template)
def main(args):
config = load_config(args.config)
check_channels(config)
check_classes(config)
palette = make_palette([classe["color"] for classe in config["classes"]])
if not args.bs:
try:
args.bs = config["model"]["bs"]
except:
pass
assert args.bs, "For rsp predict, model/bs must be set either in config file, or pass trought parameter --bs"
args.workers = args.bs if not args.workers else args.workers
cover = [tile for tile in tiles_from_csv(os.path.expanduser(args.cover))
] if args.cover else None
log = Logs(os.path.join(args.out, "log"))
if torch.cuda.is_available():
log.log("RoboSat.pink - predict on {} GPUs, with {} workers".format(
torch.cuda.device_count(), args.workers))
log.log("(Torch:{} Cuda:{} CudNN:{})".format(
torch.__version__, torch.version.cuda,
torch.backends.cudnn.version()))
device = torch.device("cuda")
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
else:
log.log("RoboSat.pink - predict on CPU, with {} workers".format(
args.workers))
log.log("")
log.log("============================================================")
log.log("WARNING: Are you -really- sure about not predicting on GPU ?")
log.log("============================================================")
log.log("")
device = torch.device("cpu")
chkpt = torch.load(args.checkpoint, map_location=device)
nn_module = load_module("robosat_pink.nn.{}".format(chkpt["nn"].lower()))
nn = getattr(nn_module, chkpt["nn"])(chkpt["shape_in"], chkpt["shape_out"],
chkpt["encoder"].lower()).to(device)
nn = torch.nn.DataParallel(nn)
nn.load_state_dict(chkpt["state_dict"])
nn.eval()
log.log("Model {} - UUID: {}".format(chkpt["nn"], chkpt["uuid"]))
with torch.no_grad(
): # don't track tensors with autograd during prediction
tiled = []
if args.passes in ["first", "both"]:
log.log("== Predict First Pass ==")
tiled = predict(config, cover, args, palette, chkpt, nn, device,
"predict")
if args.passes in ["second", "both"]:
log.log("== Predict Second Pass ==")
predict(config, cover, args, palette, chkpt, nn, device,
"predict_translate")
if not args.no_web_ui and tiled:
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, tiled, tiled, "png", template)