Function main

utils/create_strand_latent_weights.py:76–149

()

def main():

    #argparse
    parser = argparse.ArgumentParser(description='Get the weights of each dimensions after training a strand VAE')
    parser.add_argument('--checkpoint_path', required=True, help='Path to the strandVAE checkpoint')
    args = parser.parse_args()

    path_strand_vae_model=args.checkpoint_path
    

    hyperparams=HyperParamsStrandVAE() 


    normalization_dict=DiffLocksDataset.get_normalization_data()

    model = StrandCodec(do_vae=False, 
                        decode_type="dir",
                        scale_init=30.0,
                        nr_verts_per_strand=256, nr_values_to_decode=255,
                        dim_per_value_decoded=3).cuda()
    model.load_state_dict(torch.load(path_strand_vae_model))
    model = torch.compile(model)



    #latent of dimension 64 and get GT which is the mean strand
    latent=torch.zeros(1,64).cuda()
    pred_dict = model.decoder(latent, None, normalization_dict)
    pred_points=pred_dict["strand_positions"]
    gt_strand=pred_points
    print("gt_strand",gt_strand.shape)

    
    #make loss function
    loss_computer= StrandVAELoss()


    #for each dimension change it by 0.5 and check the error towards the mean strand (GT)
    loss_per_dim=[]
    for i in range(64):
        latent=torch.zeros(1,64).cuda()
        latent[:,i]=0.8
        pred_dict = model.decoder(latent, None, normalization_dict)
        pred_points=pred_dict["strand_positions"]

        #make dicts
        gt_dict={"strand_positions": gt_strand}
        pred_dict={"strand_positions": pred_points}
        latent_dict={}

        #loss
        loss_dict = loss_computer(None, gt_dict, pred_dict, latent_dict, hyperparams)
        loss=loss_dict["loss"]
        loss_per_dim.append(loss)

        # print("loss", loss)

    #normalize losses