class NaiveStridedModel(nn.Module):
"""
Baseline architecture to test out: Chain temporal models of same size:
(2d keypoints)--> [VP3D] >--(3d kypoints)--> [ModdedModel] >--(classes)
Strided version for training
"""
def __init__(self, num_joints_in, in_features, num_joints_out, filter_widths,
pretrained_weights, embedding_len, classes, causal, dropout, channels, filter_widths2=None, loadBase=True):
super().__init__()
if filter_widths2 is None:
filter_widths2 = filter_widths
self.base_model = TemporalModel(num_joints_in, in_features, num_joints_out, filter_widths,
causal=causal, dropout=dropout, channels=channels)
if loadBase:
self.base_model.load_state_dict(pretrained_weights['model_pos'])
self.top_model = ModdedStridedModel(num_joints_out, 3, num_joints_out, filter_widths2,
causal=causal, dropout=dropout, channels=embedding_len, skip_res=False)
self.top_model.shrink = nn.Conv1d( embedding_len, classes, 1)
def forward(self, x):
x = self.base_model(x)
x-= x[:,0,0,:].unsqueeze(1).unsqueeze(1)
x = self.top_model(x)
return x
def __init__(self, num_joints_in, in_features, num_joints_out, filter_widths,
pretrained_weights, embedding_len, classes, causal, dropout, channels, loadBase=True):
super().__init__()
self.base_model = TemporalModel(num_joints_in, in_features, num_joints_out, filter_widths,
causal=causal, dropout=dropout, channels=channels)
if loadBase:
self.base_model.load_state_dict(pretrained_weights['model_pos'])
self.top_model = ModdedTemporalModel(num_joints_out, 3, num_joints_out, filter_widths,
causal=causal, dropout=dropout, channels=embedding_len, skip_res=False)
self.top_model.shrink = nn.Conv1d( embedding_len, classes, 1)
def build_model(chk_filename, in_joints, in_dims, out_joints, filter_widths,
causal, channels, embedding_len, classes):
base = TemporalModel(in_joints,
in_dims,
out_joints,
filter_widths,
causal=True,
dropout=0.25,
channels=channels)
top = ModdedStridedModel(in_joints,
3,
out_joints,
filter_widths,
causal=True,
dropout=0.25,
channels=embedding_len,
skip_res=False)
base_eval = TemporalModel(in_joints,
in_dims,
out_joints,
filter_widths,
causal=True,
dropout=0.25,
channels=channels)
top_eval = ModdedTemporalModel(in_joints,
3,
out_joints,
filter_widths,
causal=True,
dropout=0.25,
channels=embedding_len,
skip_res=False)
class_mod = nn.Conv1d(embedding_len, classes, 1)
base, top, class_mod = load_model_weights(chk_filename, base, top,
class_mod)
model = nn.Sequential(
OrderedDict([('base', base),
('transform', StandardiseKeypoints(True, False)),
('embedding', HeadlessNet2(top)),
('classifier', SplitModel(class_mod))]))
eval_model = nn.Sequential(
OrderedDict([('base', base_eval),
('transform', StandardiseKeypoints(True, False)),
('embedding', HeadlessNet2(top_eval)),
('classifier', SplitModel(class_mod))]))
return model, eval_model