def save_best_model_test(self):
# create a dummy model
model = Prenet(256, out_features=[256, 256])
model = T.nn.DataParallel(layer)
# save the model
best_loss = save_best_model(model, None, 0, 100, OUT_PATH, 10, 1)
# load the model to CPU
model_dict = torch.load(MODEL_PATH,
map_location=lambda storage, loc: storage)
model.load_state_dict(model_dict['model'])
def save_checkpoint_test(self):
# create a dummy model
model = Prenet(128, out_features=[256, 128])
model = T.nn.DataParallel(layer)
# save the model
save_checkpoint(model, None, 100, OUTPATH, 1, 1)
# load the model to CPU
model_dict = torch.load(MODEL_PATH,
map_location=lambda storage, loc: storage)
model.load_state_dict(model_dict['model'])
def test_in_out(self):
layer = Prenet(128, out_features=[256, 128])
dummy_input = T.rand(4, 128)
print(layer)
output = layer(dummy_input)
assert output.shape[0] == 4
assert output.shape[1] == 128