wassname/Pointnet2_PyTorch
1
0
Fork 0
mirror of https://github.com/wassname/Pointnet2_PyTorch.git synced 2026-07-13 01:00:43 +08:00
Code Issues Packages Projects Releases Wiki Activity

Updates

Browse Source
This commit is contained in:
erikwijmans
2018-01-06 12:13:52 -05:00
parent 7e746ba72a
commit 5a5adc2b77
20 changed files with 650 additions and 494 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/pytorch_utils.py
+182 -141
View File
@@ -16,48 +16,55 @@ import math
class SharedMLP(nn.Sequential):
def __init__(self,
args: List[int],
*,
bn: bool = False,
activation=nn.ReLU(inplace=True),
name: str = ""):
def __init__(
self,
args: List[int],
*,
bn: bool = False,
activation=nn.ReLU(inplace=True),
name: str = ""
):
super().__init__()
for i in range(len(args) - 1):
self.add_module(name + 'layer{}'.format(i),
Conv2d(
args[i],
args[i + 1],
bn=bn,
activation=activation))
self.add_module(
name + 'layer{}'.format(i),
Conv2d(args[i], args[i + 1], bn=bn, activation=activation)
)
class _ConvBase(nn.Sequential):
def __init__(self,
in_size,
out_size,
kernel_size,
stride,
padding,
activation,
bn,
init,
conv=None,
batch_norm=None,
bias=True,
name=""):
def __init__(
self,
in_size,
out_size,
kernel_size,
stride,
padding,
activation,
bn,
init,
conv=None,
batch_norm=None,
bias=True,
name=""
):
super().__init__()
bias = bias and (not bn)
self.add_module(name + 'conv',
conv(
in_size,
out_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias))
self.add_module(
name + 'conv',
conv(
in_size,
out_size,
kernel_size=kernel_size,
stride=stride,
padding=padding,
bias=bias
)
)
init(self[0].weight)
if bias:
@@ -73,18 +80,21 @@ class _ConvBase(nn.Sequential):
class Conv1d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: int = 1,
stride: int = 1,
padding: int = 0,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -97,22 +107,26 @@ class Conv1d(_ConvBase):
conv=nn.Conv1d,
batch_norm=nn.BatchNorm1d,
bias=bias,
name=name)
name=name
)
class Conv2d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int] = (1, 1),
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int] = (1, 1),
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -125,22 +139,26 @@ class Conv2d(_ConvBase):
conv=nn.Conv2d,
batch_norm=nn.BatchNorm2d,
bias=bias,
name=name)
name=name
)
class Conv3d(_ConvBase):
def __init__(self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int, int] = (1, 1, 1),
stride: Tuple[int, int, int] = (1, 1, 1),
padding: Tuple[int, int, int] = (0, 0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
kernel_size: Tuple[int, int, int] = (1, 1, 1),
stride: Tuple[int, int, int] = (1, 1, 1),
padding: Tuple[int, int, int] = (0, 0, 0),
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=nn.init.kaiming_normal,
bias: bool = True,
name: str = ""
):
super().__init__(
in_size,
out_size,
@@ -153,18 +171,22 @@ class Conv3d(_ConvBase):
conv=nn.Conv3d,
batch_norm=nn.BatchNorm3d,
bias=bias,
name=name)
name=name
)
class FC(nn.Sequential):
def __init__(self,
in_size: int,
out_size: int,
*,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=None,
name: str = ""):
def __init__(
self,
in_size: int,
out_size: int,
*,
activation=nn.ReLU(inplace=True),
bn: bool = False,
init=None,
name: str = ""
):
super().__init__()
self.add_module(name + 'fc', nn.Linear(in_size, out_size, bias=not bn))
if init is not None:
@@ -183,6 +205,7 @@ class FC(nn.Sequential):
class _DropoutNoScaling(InplaceFunction):
@staticmethod
def _make_noise(input):
return input.new().resize_as_(input)
@@ -192,8 +215,9 @@ class _DropoutNoScaling(InplaceFunction):
if inplace:
return None
n = g.appendNode(
g.create("Dropout", [input]).f_("ratio", p).i_(
"is_test", not train))
g.create("Dropout", [input]).f_("ratio",
p).i_("is_test", not train)
)
real = g.appendNode(g.createSelect(n, 0))
g.appendNode(g.createSelect(n, 1))
return real
@@ -201,8 +225,10 @@ class _DropoutNoScaling(InplaceFunction):
@classmethod
def forward(cls, ctx, input, p=0.5, train=False, inplace=False):
if p < 0 or p > 1:
raise ValueError("dropout probability has to be between 0 and 1, "
"but got {}".format(p))
raise ValueError(
"dropout probability has to be between 0 and 1, "
"but got {}".format(p)
)
ctx.p = p
ctx.train = train
ctx.inplace = inplace
@@ -236,6 +262,7 @@ dropout_no_scaling = _DropoutNoScaling.apply
class _FeatureDropoutNoScaling(_DropoutNoScaling):
@staticmethod
def symbolic(input, p=0.5, train=False, inplace=False):
return None
@@ -244,7 +271,8 @@ class _FeatureDropoutNoScaling(_DropoutNoScaling):
def _make_noise(input):
return input.new().resize_(
input.size(0), input.size(1), *repeat(1,
input.dim() - 2))
input.dim() - 2)
)
feature_dropout_no_scaling = _FeatureDropoutNoScaling.apply
@@ -252,21 +280,17 @@ feature_dropout_no_scaling = _FeatureDropoutNoScaling.apply
def checkpoint_state(model=None, optimizer=None, best_prec=None, epoch=None):
return {
'epoch':
epoch,
'best_prec':
best_prec,
'model_state':
model.state_dict() if model is not None else None,
'optimizer_state':
optimizer.state_dict() if optimizer is not None else None
'epoch': epoch,
'best_prec': best_prec,
'model_state': model.state_dict() if model is not None else None,
'optimizer_state': optimizer.state_dict()
if optimizer is not None else None
}
def save_checkpoint(state,
is_best,
filename='checkpoint',
bestname='model_best'):
def save_checkpoint(
state, is_best, filename='checkpoint', bestname='model_best'
):
filename = '{}.pth.tar'.format(filename)
torch.save(state, filename)
if is_best:
@@ -325,7 +349,8 @@ def variable_size_collate(pad_val=0, use_shared_memory=True):
out = out.view(
len(batch), max_len,
*[batch[0].size(i) for i in range(1, batch[0].dim())])
*[batch[0].size(i) for i in range(1, batch[0].dim())]
)
out.fill_(pad_val)
for i in range(len(batch)):
out[i, 0:batch[i].size(0)] = batch[i]
@@ -342,8 +367,9 @@ def variable_size_collate(pad_val=0, use_shared_memory=True):
return wrapped([torch.from_numpy(b) for b in batch])
if elem.shape == (): # scalars
py_type = float if elem.dtype.name.startswith('float') else int
return _numpy_type_map[elem.dtype.name](list(
map(py_type, batch)))
return _numpy_type_map[elem.dtype.name](
list(map(py_type, batch))
)
elif isinstance(batch[0], int):
return torch.LongTensor(batch)
elif isinstance(batch[0], float):
@@ -372,19 +398,19 @@ class TrainValSplitter():
Whether or not shuffle which data goes to which split
"""
def __init__(self,
*,
numel: int,
percent_train: float,
shuffled: bool = False):
def __init__(
self, *, numel: int, percent_train: float, shuffled: bool = False
):
indicies = np.array([i for i in range(numel)])
if shuffled:
np.random.shuffle(indicies)
self.train = torch.utils.data.sampler.SubsetRandomSampler(
indicies[0:int(percent_train * numel)])
indicies[0:int(percent_train * numel)]
)
self.val = torch.utils.data.sampler.SubsetRandomSampler(
indicies[int(percent_train * numel):-1])
indicies[int(percent_train * numel):-1]
)
class CrossValSplitter():
@@ -413,7 +439,8 @@ class CrossValSplitter():
self.val = torch.utils.data.sampler.SubsetRandomSampler(self.folds[0])
self.train = torch.utils.data.sampler.SubsetRandomSampler(
np.concatenate(self.folds[1:], axis=0))
np.concatenate(self.folds[1:], axis=0)
)
self.metrics = {}
@@ -428,7 +455,8 @@ class CrossValSplitter():
assert idx >= 0 and idx < len(self)
self.val.inidicies = self.folds[idx]
self.train.inidicies = np.concatenate(
self.folds[np.arange(len(self)) != idx], axis=0)
self.folds[np.arange(len(self)) != idx], axis=0
)
def __next__(self):
self.current_v_ind += 1
@@ -454,6 +482,7 @@ class CrossValSplitter():
def set_bn_momentum_default(bn_momentum):
def fn(m):
if isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
m.momentum = bn_momentum
@@ -462,14 +491,17 @@ def set_bn_momentum_default(bn_momentum):
class BNMomentumScheduler(object):
def __init__(self,
model,
bn_lambda,
last_epoch=-1,
setter=set_bn_momentum_default):
def __init__(
self, model, bn_lambda, last_epoch=-1,
setter=set_bn_momentum_default
):
if not isinstance(model, nn.Module):
raise RuntimeError("Class '{}' is not a PyTorch nn Module".format(
type(model).__name__))
raise RuntimeError(
"Class '{}' is not a PyTorch nn Module".format(
type(model).__name__
)
)
self.model = model
self.setter = setter
@@ -511,18 +543,21 @@ class Trainer(object):
Name of file to output tensorboard_logger to
"""
def __init__(self,
model,
model_fn,
optimizer,
checkpoint_name="ckpt",
best_name="best",
lr_scheduler=None,
bnm_scheduler=None,
eval_frequency=1,
log_name=None):
def __init__(
self,
model,
model_fn,
optimizer,
checkpoint_name="ckpt",
best_name="best",
lr_scheduler=None,
bnm_scheduler=None,
eval_frequency=1,
log_name=None
):
self.model, self.model_fn, self.optimizer, self.lr_scheduler, self.bnm_scheduler = (
model, model_fn, optimizer, lr_scheduler, bnm_scheduler)
model, model_fn, optimizer, lr_scheduler, bnm_scheduler
)
self.checkpoint_name, self.best_name = checkpoint_name, best_name
self.eval_frequency = eval_frequency
@@ -536,7 +571,8 @@ class Trainer(object):
@staticmethod
def _print(mode, epoch, loss, eval_dict, count):
to_print = "[{:d}] {}\tMean Loss: {:.4e}".format(
epoch, mode, loss / count)
epoch, mode, loss / count
)
for k, v in natsorted(eval_dict.items(), key=itemgetter(0)):
to_print += "\tMean {}: {:2.3f}%".format(k, stats.mean(v) * 1e2)
@@ -574,7 +610,8 @@ class Trainer(object):
for k, v in eval_res.items():
if v is not None:
tb_log.log_value(
"Training {}".format(k), 1.0 - v, step=idx)
"Training {}".format(k), 1.0 - v, step=idx
)
d_loader.dataset.randomize()
@@ -593,7 +630,8 @@ class Trainer(object):
self.optimizer.zero_grad()
_, loss, eval_res = self.model_fn(
self.model, data, eval=True, epoch=epoch)
self.model, data, eval=True, epoch=epoch
)
total_loss += loss.data[0]
count += 1
@@ -606,8 +644,7 @@ class Trainer(object):
tb_log.log_value("Eval loss", loss.data[0], step=idx)
for k, v in eval_res.items():
if v is not None:
tb_log.log_value(
"Eval {}".format(k), 1.0 - v, step=idx)
tb_log.log_value("Eval {}".format(k), 1.0 - v, step=idx)
d_loader.dataset.randomize()
@@ -615,12 +652,14 @@ class Trainer(object):
return total_loss / count, eval_dict
def train(self,
start_epoch,
n_epochs,
train_loader,
test_loader=None,
best_loss=0.0):
def train(
self,
start_epoch,
n_epochs,
train_loader,
test_loader=None,
best_loss=0.0
):
r"""
Call to begin training the model
@@ -649,10 +688,12 @@ class Trainer(object):
is_best = val_loss < best_loss
best_loss = min(best_loss, val_loss)
save_checkpoint(
checkpoint_state(self.model, self.optimizer, val_loss,
epoch),
checkpoint_state(
self.model, self.optimizer, val_loss, epoch
),
is_best,
filename=self.checkpoint_name,
bestname=self.best_name)
bestname=self.best_name
)
return best_loss