RST/evaluate_ethucy.py at main · cimat-ris/RST · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
from utils.datasets_utils import get_dataset
from utils.config import load_config
from utils.constants import *
from utils.metrics import min_ade,min_fde,print_results_ethucy
import torch
import numpy as np
from models.generate import generate_ddim,generate_shift
from models.unet import TrajectoryDenoiser_CondEmbed,TrajectoryDenoiser_CondMerge,TrajectoryDenoiser_Shift
from models.training import train_one_epoch,train_one_epoch_shift
from models.diffusion.ddpm import DDPM
from models.diffusion.resshift import ResShift
from pathlib import Path
import numpy as np
import gc,logging, os
from tqdm import tqdm

logging.basicConfig(format='%(levelname)s: %(message)s',level=20)

# Device to use later on
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info("Using device: "+torch.cuda.get_device_name(DEVICE))

# Load configuation file (conditional model)
config = load_config("ethucy-conditional-past.yaml")
print(config)
#config = load_config("ethucy-conditional-past-shift.yaml")
# config = load_config("ethucy-conditional-past-social.yaml")
# Load configuation file (unconditional model)
# config = load_config("ethucy-unconditional.yaml")
all_mades          = []
all_mfdes          = []
mades_per_sequence = []
mfdes_per_sequence = []

# Cycle over the possible test datasets in ETH-UCY
for id_test in range(5):
    config["dataset"]["id_test"] = id_test
    logging.info("Evaluation on: "+SUBDATASETS_NAMES[0][id_test])
    # Get the data
    batched_train_data,batched_val_data,batched_test_data,homography,reference_image = get_dataset(config["dataset"])

    # Instantiate the denoiser
    if config["diffusion"]["sampler"] == "shift":
        denoiser = TrajectoryDenoiser_Shift(config["model"])
    else:
        if config["model"]["condition_handling"] == "merge":
            denoiser   = TrajectoryDenoiser_CondMerge(config["model"])
        else:
            denoiser   = TrajectoryDenoiser_CondEmbed(config["model"])
    denoiser.to(DEVICE)

    # The optimizer (Adam with weight decay)
    optimizer= torch.optim.AdamW(denoiser.parameters(),lr=config["train"]["initial_lr"])

    # Instantiate the diffusion model
    if config["diffusion"]["sampler"] == "shift":
        diffusionmodel = ResShift(timesteps=config["diffusion"]["timesteps"],kappa=config["diffusion"]["kappa"])
    else:
        diffusionmodel = DDPM(timesteps=config["diffusion"]["timesteps"])
    diffusionmodel.to(DEVICE)

    # Training loop
    best_loss = 1e6
    for epoch in range(1,config["train"]["epochs"] + 1):
        torch.cuda.empty_cache()
        gc.collect()

        # One epoch of training
        if config["diffusion"]["sampler"] == "shift":
            epoch_loss= train_one_epoch_shift(denoiser,diffusionmodel,batched_train_data,optimizer,DEVICE,epoch=epoch,total_epochs=config["train"]["epochs"])
        else:
            epoch_loss= train_one_epoch(denoiser,diffusionmodel,batched_train_data,optimizer,DEVICE,epoch=epoch,total_epochs=config["train"]["epochs"])
        if epoch_loss < best_loss:
            best_loss = epoch_loss
            # Save best checkpoints
            checkpoint_dict = {
                "opt": optimizer.state_dict(),
                "model": denoiser.state_dict()
            }
            if not os.path.exists(config["model"]["save_dir"]):
		        # Create a new directory if it does not exist
                os.makedirs(config["model"]["save_dir"])
            save_path = config["model"]["save_dir"]+(config["model"]["model_name"].format(config["model"]["condition"],config["train"]["epochs"]))
            torch.save(checkpoint_dict, save_path)
            del checkpoint_dict

    # Evaluate metrics on the best model
    save_path = config["model"]["save_dir"]+(config["model"]["model_name"].format(config["model"]["condition"],config["train"]["epochs"]))
    denoiser.load_state_dict(torch.load(save_path)["model"])
    mades = []
    mfdes = []
    step  = config["diffusion"]["trajs_at_a_time"]
    for batch in tqdm(iterable=batched_test_data,dynamic_ncols=False,total=len(batched_test_data),desc="Batches :: ", position=0):
        past_velocities_batch,__,past_positions_batch,future_positions_batch,neighbors_batch,__,neighbors_mask_batch = batch
        for id_trajectory in range(0,past_velocities_batch.shape[0],step):
            # Past and future absolute positions
            past_positions      = past_positions_batch[id_trajectory:id_trajectory+step,:,:]
            future_positions    = future_positions_batch[id_trajectory:id_trajectory+step,:,:]
            # Past velocities (used as a condition)
            past_velocities     = past_velocities_batch[id_trajectory:id_trajectory+step,:,:]
            past_velocities     = past_velocities.to(DEVICE)
            # Neighbors
            if neighbors_batch is None:
                neighbors       = None
                neighbors_mask  = None
            else:
                neighbors       = neighbors_batch[id_trajectory:id_trajectory+step,:,:,:]
                neighbors_mask  = neighbors_mask_batch[id_trajectory:id_trajectory+step,:,:]
                neighbors       = neighbors.to(DEVICE)
                neighbors_mask  = neighbors_mask.to(DEVICE)
            if config["diffusion"]["sampler"] == "shift":
                future_rough  = torch.zeros_like(future_positions)
                future_rough[:,:,:]= past_velocities[:,-2:-1,:]
                future_rough  = future_rough.to(DEVICE)
                predicted_velocities= generate_shift(denoiser,backward_sampler=diffusionmodel,past=past_velocities,rough=future_rough,config=config,device=DEVICE).cpu()
            else:
                taus  = np.arange(0,config["diffusion"]["timesteps"],config["diffusion"]["ddim_divider"])
                predicted_velocities= generate_ddim(denoiser,taus,diffusionmodel,past_velocities,neighbors,neighbors_mask,config,DEVICE,False).cpu()
            predicted_positions = 0.4*np.cumsum(predicted_velocities[:,:,:].permute(0,2,1),axis=1)
            predicted_positions = predicted_positions+torch.repeat_interleave(past_positions, repeats=config["diffusion"]["nsamples"], dim=0)[:,-1:,:]
            # Evaluade made and mdfe
            for i in range(step):
                nsamples = config["diffusion"]["nsamples"]
                if i*nsamples>=predicted_positions.shape[0]:
                    break
                made     = min_ade(predicted_positions[i*nsamples:(i+1)*nsamples,:,:],future_positions[i,:,:])
                mfde     = min_fde(predicted_positions[i*nsamples:(i+1)*nsamples,:,:],future_positions[i,:,:])
                mades.append(made)
                mfdes.append(mfde)
                all_mades.append(made)
                all_mfdes.append(mfde)

    mades_per_sequence.append(np.mean(np.array(mades)))
    mfdes_per_sequence.append(np.mean(np.array(mfdes)))
    logging.info('Sequence: {}  MADE: {:02.4f} MFDE: {:02.4f}'.format(SUBDATASETS_NAMES[0][id_test],mades_per_sequence[id_test],mfdes_per_sequence[id_test]))

# Print results
print_results_ethucy(mades_per_sequence,mfdes_per_sequence,all_mades,all_mfdes,sota_comparison=True)