def get_transformations(transform_spec):
transformer = transformers.TransformerSequence(*[
TRANSFORMER_STRING_MAP[trans_name](networks.EquivariantPosePredictor,
1, 32)
for trans_name in transform_spec
])
transforms = T.TransformSequence(
*[TRANSFORM_STRING_MAP[trans_name]() for trans_name in transform_spec])
return transforms, transformer
writer.close()
if __name__ == "__main__":
csv = "gz2_data/gz_amended.csv"
img = "gz2_data/"
a01 = "t01_smooth_or_features_a01_smooth_count"
a02 = "t01_smooth_or_features_a02_features_or_disk_count"
a03 = "t01_smooth_or_features_a03_star_or_artifact_count"
list_of_ans = [a01, a02, a03]
data = Gz2_data(csv_dir=csv,
image_dir=img,
list_of_interest=list_of_ans,
crop=80,
resize=80)
trans = transformers.TransformerSequence(
transformers.Translation(networks.EquivariantPosePredictor, 1, 32),
transformers.Rotation(networks.EquivariantPosePredictor, 1, 32))
encoder_args = {'insize': 80, 'z_dim': 100, 'transformer': trans}
decoder_args = {'z_dim': 100, 'outsize': 80}
test_proportion = 0.1
train_loader, test_loader = return_data_loader(data, test_proportion, batch_size=2, shuffle=True)
optimizer = Adam({"lr": 1e-4})
pvae = PoseVAE(Encoder, Decoder, 10, encoder_args, decoder_args, use_cuda=False)
svi = SVI(pvae.model, pvae.guide, optimizer, loss=Trace_ELBO())
train_log("test/", pvae, svi, train_loader, test_loader, 1, use_cuda=False)
