RealUID/generate_samples.py at main · David-cripto/RealUID · 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
import argparse
import random
import torch
try:
    import torch_npu
    from torch_npu.contrib import transfer_to_npu
except Exception:
    print("Not npu case")
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
print(device)


import os
import numpy as np


from src.models import UNetModelWrapperWithHead
from src.generate import  generate_and_save_samples


parser = argparse.ArgumentParser(description="model parser")


#paths
parser.add_argument('--output_dir', type=str) #"./result_cifar/"
parser.add_argument('--exp_name', type=str, default=None, help='name of the subfolder in the output dir')

parser.add_argument('--model_ckpt_path', type=str, help = 'model checkpoint')
parser.add_argument('--mode',default='one_step',  choices=['one_step', 'multi_step'], help='one-step mode or flow vector field')
parser.add_argument('--ema_decay', type=float, default=None, help='ema decay which to evaluate if given. if not given the non-ema net is evaluated.')
#dataset
parser.add_argument('--dataset', default='cifar10',  choices=['cifar10', 'celeba'])
parser.add_argument('--cond', default='uncond',  choices=['cond', 'uncond'])
parser.add_argument('--seed', type=int, default=42)
parser.add_argument('--num_classes', type=int, default=10)


parser.add_argument('--num_save_image', type=int, default=100, help='number of samples to generate')

args = parser.parse_args()


COND = args.cond == 'cond' #conditional or unconditional


num_save_image = args.num_save_image
dataset_name = args.dataset
one_step = args.mode == 'one_step'
ema_decay = args.ema_decay

#Set seed
seed = args.seed
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # for multi-GPU
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
np.random.seed(seed)
random.seed(seed)


exp_name = args.exp_name
if exp_name is None:
    exp_name = f'{args.dataset}_{args.cond}_generated_samples'

output_dir = args.output_dir


savedir = os.path.join(output_dir, exp_name)
os.makedirs(savedir, exist_ok=True)


###Init allmodels

resolution = 32 if dataset_name == 'cifar10' else 64
num_classes = args.num_classes
model = UNetModelWrapperWithHead(
        dim=(3, resolution, resolution),
        num_res_blocks=2,
        num_channels=128,
        channel_mult=[1, 2, 2, 2],
        num_heads=4,
        num_head_channels=64,
        attention_resolutions="16",
        dropout=0.0,
        class_cond = COND,
        num_classes = num_classes
    ).to(device)


print("path: ", args.model_ckpt_path)
checkpoint = torch.load(args.model_ckpt_path, map_location=device)

if one_step:
    if ema_decay is None:
        model.load_state_dict(checkpoint["gen"], strict = False)
    else:
        model.load_state_dict(checkpoint["ema_gens"][str(ema_decay)], strict = False)
else:
    if ema_decay is None:
        model.load_state_dict(checkpoint["net_model"], strict = False)
    else:
        model.load_state_dict(checkpoint["ema_models"][str(ema_decay)], strict = False)


model.eval()

fids = []


with torch.no_grad():

    if COND:
        y = torch.arange(num_save_image, device=device, dtype=int) % num_classes
    else:
        y = None

    generate_and_save_samples(model, savedir, y = y, one_step=one_step, name ="generated_samples", batch_size=num_save_image, step = 0)