def __init__(self, config):
super(KPFCNN, self).__init__()
self.encoder = KPCNN(config)
self.config = config
self.blocks = nn.ModuleDict()
# Feature Extraction Module
# Find first upsampling block
start_i = 0
for block_i, block in enumerate(config.architecture):
if 'upsample' in block:
start_i = block_i
break
layer = config.num_layers - 1
r = config.first_subsampling_dl * config.density_parameter * 2 ** layer
in_fdim = config.first_features_dim * 2 ** layer
out_fdim = in_fdim
block_in_layer = 0
for block_i, block in enumerate(config.architecture[start_i:]):
is_strided = 'strided' in block
self.blocks[f'layer{layer}/{block}'] = get_block(block, config, int(1.5 * in_fdim), out_fdim, radius=r, strided=is_strided)
# update feature dimension
in_fdim = out_fdim
block_in_layer += 1
# Detect change to a subsampled layer
if 'upsample' in block:
# Update radius and feature dimension for next layer
out_fdim = out_fdim // 2
r *= 0.5
layer -= 1
block_in_layer = 0
# Segmentation Head
self.blocks['segmentation_head'] = nn.Sequential(
nn.Linear(out_fdim, config.first_features_dim),
nn.BatchNorm1d(config.first_features_dim, momentum=config.batch_norm_momentum, eps=1e-6),
nn.LeakyReLU(negative_slope=0.2),
# nn.Dropout(p=config.dropout),
nn.Linear(config.first_features_dim, config.num_classes)
)
def __init__(self, config):
is_test = False
if is_test:
self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M') + 'Test'
else:
self.experiment_id = "KPConvNet" + time.strftime('%m%d%H%M')
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.verbose = True
# snapshot
self.snapshot_interval = 5
snapshot_root = f'snapshot/{config.dataset}_{self.experiment_id}'
tensorboard_root = f'tensorboard/{config.dataset}_{self.experiment_id}'
os.makedirs(snapshot_root, exist_ok=True)
os.makedirs(tensorboard_root, exist_ok=True)
shutil.copy2(os.path.join('.', 'training_ShapeNetCls.py'), os.path.join(snapshot_root, 'train.py'))
shutil.copy2(os.path.join('datasets', 'ShapeNet.py'), os.path.join(snapshot_root, 'dataset.py'))
shutil.copy2(os.path.join('datasets', 'dataloader.py'), os.path.join(snapshot_root, 'dataloader.py'))
self.save_dir = os.path.join(snapshot_root, 'models/')
self.result_dir = os.path.join(snapshot_root, 'results/')
self.tboard_dir = tensorboard_root
# dataset & dataloader
self.train_set = ShapeNetDataset(root=config.data_train_dir,
split='train',
first_subsampling_dl=config.first_subsampling_dl,
classification=True,
config=config,
)
self.test_set = ShapeNetDataset(root=config.data_test_dir,
split='test',
first_subsampling_dl=config.first_subsampling_dl,
classification=True,
config=config,
)
self.train_loader = get_dataloader(dataset=self.train_set,
batch_size=config.train_batch_size,
shuffle=True,
num_workers=4,
)
self.test_loader = get_dataloader(dataset=self.test_set,
batch_size=config.test_batch_size,
shuffle=False,
num_workers=4,
)
print("Training set size:", self.train_loader.dataset.__len__())
print("Test set size:", self.test_loader.dataset.__len__())
# model
self.model = KPCNN(config)
self.resume = config.resume
# optimizer
self.start_epoch = 0
self.epoch = config.max_epoch
self.optimizer = optim.SGD(self.model.parameters(), lr=config.learning_rate, momentum=config.momentum, weight_decay=1e-6)
self.scheduler = optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=config.exp_gamma)
self.scheduler_interval = config.exp_interval
# evaluate
self.evaluate_interval = 1
self.evaluate_metric = nn.CrossEntropyLoss(reduction='mean')
self.check_args()
class KPFCNN(nn.Module):
def __init__(self, config):
super(KPFCNN, self).__init__()
self.encoder = KPCNN(config)
self.config = config
self.blocks = nn.ModuleDict()
# Feature Extraction Module
# Find first upsampling block
start_i = 0
for block_i, block in enumerate(config.architecture):
if 'upsample' in block:
start_i = block_i
break
layer = config.num_layers - 1
r = config.first_subsampling_dl * config.density_parameter * 2 ** layer
in_fdim = config.first_features_dim * 2 ** layer
out_fdim = in_fdim
block_in_layer = 0
for block_i, block in enumerate(config.architecture[start_i:]):
is_strided = 'strided' in block
self.blocks[f'layer{layer}/{block}'] = get_block(block, config, int(1.5 * in_fdim), out_fdim, radius=r, strided=is_strided)
# update feature dimension
in_fdim = out_fdim
block_in_layer += 1
# Detect change to a subsampled layer
if 'upsample' in block:
# Update radius and feature dimension for next layer
out_fdim = out_fdim // 2
r *= 0.5
layer -= 1
block_in_layer = 0
# Segmentation Head
self.blocks['segmentation_head'] = nn.Sequential(
nn.Linear(out_fdim, config.first_features_dim),
nn.BatchNorm1d(config.first_features_dim, momentum=config.batch_norm_momentum, eps=1e-6),
nn.LeakyReLU(negative_slope=0.2),
# nn.Dropout(p=config.dropout),
nn.Linear(config.first_features_dim, config.num_classes)
)
# print(list(self.named_parameters()))
def forward(self, inputs):
features = self.feature_extraction(inputs)
logits = self.segmentation_head(features)
return logits
def feature_extraction(self, inputs):
F = self.encoder.feature_extraction(inputs)
features = F[-1]
# Current radius of convolution and feature dimension
layer = self.config.num_layers - 1
r = self.config.first_subsampling_dl * self.config.density_parameter * 2 ** layer
fdim = self.config.first_features_dim * 2 ** layer
# Find first upsampling block
start_i = 0
for block_i, block in enumerate(self.config.architecture):
if 'upsample' in block:
start_i = block_i
break
# Loop over upsampling blocks
for block_i, block in enumerate(self.config.architecture[start_i:]):
# Get the function for this layer
block_ops = self.blocks[f'layer{layer}/{block}']
# Apply the layer function defining tf ops
if 'upsample' in block:
if block == 'nearest_upsample':
upsample_indices = inputs['upsamples'][layer - 1]
else:
raise ValueError(f"Unknown block type. {block}")
features = block_ops(upsample_indices, features)
else:
if block in ['unary', 'simple', 'resnet', 'resnetb']:
query_points = inputs['points'][layer]
support_points = inputs['points'][layer]
neighbors_indices = inputs['neighbors'][layer]
elif block in ['simple_strided', 'resnetb_strided', 'resnetb_deformable_strided']:
query_points = inputs['points'][layer + 1]
support_points = inputs['points'][layer]
neighbors_indices = inputs['pools'][layer]
else:
raise ValueError(f"Unknown block type. {block}")
features = block_ops(query_points, support_points, neighbors_indices, features)
# Detect change to a subsampled layer
if 'upsample' in block:
# Update radius and feature dimension for next layer
layer -= 1
r *= 0.5
fdim = fdim // 2
# Concatenate with CNN feature map
features = torch.cat((features, F[layer]), dim=1)
return features
def segmentation_head(self, features):
logits = self.blocks['segmentation_head'](features)
return logits
def train():
logger.info("test Dataloader")
cfg = get_cfg_defaults()
cfg.merge_from_file("./yaml/modelnet.yaml")
writer = SummaryWriter()
v = cfg.NETWORK.FIRST_SUBSAMPLING_DL
mode_neigh = 0 # neighbors or edge
architecture, list_voxel_size, list_radius, \
list_radius_conv, list_size = get_list_constants(
cfg.NETWORK.FIRST_SUBSAMPLING_DL, cfg.NETWORK.DENSITY_PARAMETER/2,
cfg.NETWORK.ARCHITECTURE, cfg.NETWORK.FIRST_DIM,
cfg.INPUT.IN_FEATURES_DIM,
cfg.INPUT.NUM_CLASSES)
list_transfo = [DataSubsampling(list_voxel_size)]
neigh = Neighbors(list_radius,
max_num_neighbors=cfg.INPUT.MAX_NUM_NEIGHBORS,
is_pool=True,
is_upsample=False,
mode=mode_neigh)
list_transfo.append(neigh)
transfo = Compose(list_transfo)
# transfo = DataSubsampling(list_voxel_size)
dataset = ModelNet("./Data/ModelNet40",
subsampling_parameter=v,
transforms=transfo)
dataloader = MultiScaleDataLoader(dataset,
batch_size=cfg.NETWORK.BATCH_NUM,
shuffle=True,
num_workers=cfg.SYSTEM.NUM_WORKERS,
pin_memory=False)
model = KPCNN(architecture,
list_radius_conv,
list_size,
cfg,
mode=mode_neigh)
optimizer = torch.optim.SGD(model.parameters(),
lr=cfg.TRAIN.LEARNING_RATE,
momentum=cfg.TRAIN.MOMENTUM)
model.cuda()
for e in range(1):
logger.info("epoch {:d}", e)
model.train()
for i, batch in enumerate(dataloader):
if (mode_neigh == 0):
batch = shadow_neigh(batch).to('cuda')
else:
batch = batch.to('cuda')
# print(batch.list_neigh[3][:20])
pred = model(batch)
loss = compute_classification_loss(pred, batch.y, model,
cfg.TRAIN.WEIGHTS_DECAY)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (i % cfg.TRAIN.LOG_INTERVAL == 0):
accuracy = compute_classification_accuracy(pred, batch.y)
global_step = e * len(dataloader) + i
logger.info("Epoch: {} Step {}, loss {:3f}, accuracy: {:3f}",
e, i, loss.item(), accuracy.item())
writer.add_scalar('Loss/train', loss.item(), global_step)
writer.add_scalar('Accuracy/train', accuracy.item(),
global_step)