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"))
type=int,
help="number of superpixels")
parser.add_argument("--pos_scale", "-p", default=10, type=float)
parser.add_argument("--folder",
"-f",
default='log',
help="a folder to store result")
args = parser.parse_args()
if not os.path.exists(args.folder):
os.mkdir(args.folder)
data = shapenet("../../shapenet_part_seg_hdf5_data", split='val')
loader = DataLoader(data, batch_size=1, shuffle=False)
model = PointNet_SSN(args.fdim,
args.nspix,
args.niter,
backend=soft_slic_pknn).to("cuda")
model.load_state_dict(torch.load(args.weight))
model.eval()
print(model)
s = time.time()
for i, (pointcloud, label, labell) in enumerate(loader):
print(i)
_, labels, _, _ = inference(pointcloud, 10, model)
pointcloud = pointcloud.squeeze(0).transpose(1, 0).numpy()
label = labell.transpose(1, 0).numpy()
#spix = spix.squeeze(0).transpose(1, 0).numpy()
ptcloud = np.concatenate(
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)
train_dataset = shapenet.shapenet(cfg.root)
train_loader = DataLoader(train_dataset,
cfg.batchsize,
shuffle=True,
drop_last=True,
num_workers=cfg.nworkers)
test_dataset = shapenet.shapenet(cfg.root, split="test")
test_loader = DataLoader(test_dataset, 1, shuffle=False, drop_last=False)
meter = Meter()
iterations = 0
max_val_asa = 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(
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("lr", metric["lr"], iterations)
if (iterations % 1000) == 0:
torch.save(
model.state_dict(),
os.path.join(
cfg.out_dir, "model_epoch_" + str(epoch_idx) + "_" +
str(batch_iterations) + '_iter_' + str(iterations) +
".pth"))
# if (iterations % cfg.test_interval) == 0:
# asa = eval(model, test_loader, cfg.pos_scale, device)
# print(f"validation asa {asa}")
# writer.add_scalar("comprehensive/asa", asa, iterations)
# if asa > max_val_asa:
# max_val_asa = asa
# torch.save(model.state_dict(), os.path.join(
# cfg.out_dir, "bset_model_sp_loss.pth"))
unique_id = str(int(time.time()))
torch.save(model.state_dict(),
os.path.join(cfg.out_dir, "model" + unique_id + ".pth"))