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"))
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)
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"))
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(