示例#1
0
def train(cfg):
    if torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"

    model = PointNet_SSN(cfg.fdim,
                         cfg.nspix,
                         cfg.niter,
                         backend=soft_slic_pknn).to(device)

    optimizer = optim.Adam(model.parameters(), cfg.lr)
    decayer = optim.lr_scheduler.StepLR(optimizer, 1, 0.95)

    train_dataset = shapenet.shapenet_inst(cfg.root)
    train_loader = DataLoader(train_dataset,
                              cfg.batchsize,
                              shuffle=True,
                              drop_last=True,
                              num_workers=cfg.nworkers)

    test_dataset = shapenet.shapenet_inst(cfg.root, split="test")
    test_loader = DataLoader(test_dataset, 1, shuffle=False, drop_last=False)

    meter = Meter()

    iterations = 0
    writer = SummaryWriter(log_dir=args.out_dir, comment='traininglog')
    for epoch_idx in range(cfg.train_epoch):
        batch_iterations = 0
        for data in train_loader:
            batch_iterations += 1
            iterations += 1
            metric = update_param(data, model, optimizer, cfg.compactness,
                                  cfg.pos_scale, device)
            meter.add(metric)
            state = meter.state(
                f"[{batch_iterations},{epoch_idx}/{cfg.train_epoch}]")
            print(state)
            addscaler(metric, writer, iterations)
            if (iterations % 200) == 0:
                test_res = eval(model, test_loader, cfg.pos_scale, device)
                addscaler(test_res, writer, iterations, True)
                if (iterations % 1000) == 0:
                    (asa, usa) = test_res
                    strs = "ep_{:}_batch_{:}_iter_{:}_asa_{:.3f}_ue_{:.3f}.pth".format(
                        epoch_idx, batch_iterations, iterations, asa, usa)
                    torch.save(model.state_dict(),
                               os.path.join(args.out_dir, strs))
        decayer.step()
    unique_id = str(int(time.time()))
    torch.save(model.state_dict(),
               os.path.join(cfg.out_dir, "model" + unique_id + ".pth"))

    unique_id = str(int(time.time()))
    torch.save(model.state_dict(),
               os.path.join(cfg.out_dir, "model" + unique_id + ".pth"))
示例#2
0
def train(cfg):
    if torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"

    model = MfaPtnetSsn(cfg.fdim, cfg.nspix, cfg.niter,
                        backend=soft_slic_pknn).to(device)

    disc_loss = discriminative_loss(0.1, 0.5)

    optimizer = optim.Adam(model.parameters(), cfg.lr)
    dacayer = optim.lr_scheduler.StepLR(optimizer, 2, 0.9)

    train_dataset = shapenet.shapenet_inst(cfg.root)
    train_loader = DataLoader(train_dataset,
                              cfg.batchsize,
                              shuffle=True,
                              drop_last=True,
                              num_workers=cfg.nworkers)
示例#3
0
def train(cfg):
    if torch.cuda.is_available():
        device = "cuda"
    else:
        device = "cpu"

    model = bistream_SSN(cfg.fdim,
                         cfg.nspix,
                         cfg.niter,
                         backend=soft_slic_pknn).to(device)

    disc_loss = discriminative_loss(0.1, 0.5, 1e-4)
    optimizer = optim.Adam(model.parameters(), cfg.lr)
    dacayer = optim.lr_scheduler.StepLR(optimizer, 2, 0.9)
    train_dataset = shapenet.shapenet_inst(cfg.root)

    train_loader = DataLoader(train_dataset,
                              cfg.batchsize,
                              shuffle=True,
                              drop_last=True,
                              num_workers=cfg.nworkers)

    test_dataset = shapenet.shapenet_inst(cfg.root, split="test")
    test_loader = DataLoader(test_dataset, 1, shuffle=False, drop_last=False)

    meter = Meter()

    iterations = 0
    writer = SummaryWriter(log_dir=cfg.out_dir, comment='traininglog')
    for epoch_idx in range(cfg.train_epoch):
        batch_iterations = 0
        for data in train_loader:
            batch_iterations += 1
            iterations += 1
            metric = update_param(data, model, optimizer, cfg.compactness,
                                  cfg.pos_scale, device, disc_loss)
            meter.add(metric)
            state = meter.state(
                f"[{batch_iterations},{epoch_idx}/{cfg.train_epoch}]")
            print(state)
            # return {"loss": loss.item(), "reconstruction": recons_loss.item(), "compact": compact_loss.item()}
            writer.add_scalar("comprehensive/loss", metric["loss"], iterations)
            writer.add_scalar("loss/reconstruction_loss",
                              metric["reconstruction"], iterations)
            writer.add_scalar("loss/compact_loss", metric["compact"],
                              iterations)
            writer.add_scalar("loss/disc_loss", metric["disc_loss"],
                              iterations)
            writer.add_scalar("lr",
                              
                              optimizer.state_dict()['param_groups'][0]['lr'],
                              
                              iterations)
            if (iterations % 200) == 0:
                (asa, usa) = eval(model, test_loader, cfg.pos_scale, device)
                writer.add_scalar("test/asa", asa, iterations)
                writer.add_scalar("test/ue", usa, iterations)
                if (iterations % 1000) == 0:
                    strs = "ep_{:}_batch_{:}_iter_{:}_asa_{:.3f}_ue_{:.3f}.pth".format(
                        epoch_idx, batch_iterations, iterations, asa, usa)
                    torch.save(model.state_dict(),
                               os.path.join(cfg.out_dir, strs))
        dacayer.step()

    unique_id = str(int(time.time()))
    torch.save(model.state_dict(),
               os.path.join(cfg.out_dir, "model" + unique_id + ".pth"))
示例#4
0
    model = MfaPtnetSsn(cfg.fdim, cfg.nspix, cfg.niter,
                        backend=soft_slic_pknn).to(device)

    disc_loss = discriminative_loss(0.1, 0.5)

    optimizer = optim.Adam(model.parameters(), cfg.lr)
    dacayer = optim.lr_scheduler.StepLR(optimizer, 2, 0.9)

    train_dataset = shapenet.shapenet_inst(cfg.root)
    train_loader = DataLoader(train_dataset,
                              cfg.batchsize,
                              shuffle=True,
                              drop_last=True,
                              num_workers=cfg.nworkers)
.
    test_dataset = shapenet.shapenet_inst(cfg.root, split="test")
    test_loader = DataLoader(test_dataset, 1, shuffle=False, drop_last=False)

    meter = Meter()

    iterations = 0
    writer = SummaryWriter(log_dir=cfg.out_dir, comment='traininglog')
    for epoch_idx in range(cfg.train_epoch):
        batch_iterations = 0
        for data in train_loader:
            batch_iterations += 1
            iterations += 1
            metric = update_param(data, model, optimizer, cfg.compactness,
                                  cfg.pos_scale, device)
            meter.add(metric)
            state = meter.state(