File size: 5,214 Bytes
28e6f98 | 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 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | from diffusion_pytorch import GaussianDiffusion, Trainer, Model
from Fid import calculate_fid_given_samples
import torchvision
import os
import errno
import shutil
import argparse
def create_folder(path):
try:
os.mkdir(path)
except OSError as exc:
if exc.errno != errno.EEXIST:
raise
pass
def del_folder(path):
try:
shutil.rmtree(path)
except OSError as exc:
pass
create = 0
if create:
trainset = torchvision.datasets.CIFAR10(
root='./data', train=False, download=True)
root = './root_cifar10_test/'
del_folder(root)
create_folder(root)
for i in range(10):
lable_root = root + str(i) + '/'
create_folder(lable_root)
for idx in range(len(trainset)):
img, label = trainset[idx]
print(idx)
img.save(root + str(label) + '/' + str(idx) + '.png')
parser = argparse.ArgumentParser()
parser.add_argument('--time_steps', default=50, type=int)
parser.add_argument('--sample_steps', default=None, type=int)
parser.add_argument('--kernel_std', default=0.1, type=float)
parser.add_argument('--save_folder', default='progression_cifar', type=str)
parser.add_argument('--load_path', default='/cmlscratch/eborgnia/cold_diffusion/paper_defading_random_1/model.pt', type=str)
parser.add_argument('--data_path', default='./root_cifar10_test/', type=str)
parser.add_argument('--test_type', default='test_paper_showing_diffusion_images_diff', type=str)
parser.add_argument('--fade_routine', default='Random_Incremental', type=str)
parser.add_argument('--sampling_routine', default='x0_step_down', type=str)
parser.add_argument('--remove_time_embed', action="store_true")
parser.add_argument('--discrete', action="store_true")
parser.add_argument('--residual', action="store_true")
args = parser.parse_args()
print(args)
img_path=None
if 'train' in args.test_type:
img_path = args.data_path
elif 'test' in args.test_type:
img_path = args.data_path
print("Img Path is ", img_path)
image_channels = 1
if model_name == "unet":
model = Model(resolution=args.image_size,
in_channels=1,
out_ch=1,
ch=128,
ch_mult=(1, 2, 2, 2),
num_res_blocks=2,
attn_resolutions=(16,),
dropout=0.1).cuda()
elif model_name == "twounet":
model = TwoBranchNewModel(resolution=args.image_size,
in_channels=1,
out_ch=1,
ch=128,
ch_mult=(1, 2, 2, 2),
num_res_blocks=2,
attn_resolutions=(16,),
dropout=0.1).cuda()
elif model_name == "twobranch":
downsample = [4, 4, 4]
disc_channels = 64
disc_layers = 3
discriminator_iter_start = 10000
disc_loss_type = "hinge"
image_gan_weight = 1.0
video_gan_weight = 1.0
l1_weight = 4.0
gan_feat_weight = 4.0
perceptual_weight = 4.0
i3d_feat = False
restart_thres = 1.0
no_random_restart = False
norm_type = "group"
padding_type = "replicate"
num_groups = 32
base_num_every_group = 2
num_features = 64
act = "PReLU"
num_channels = 1
model = TwoBranchModel(
image_channels,
disc_channels, disc_layers, disc_loss_type,
gan_feat_weight, image_gan_weight,
discriminator_iter_start,
perceptual_weight, l1_weight,
num_features, act, base_num_every_group, num_channels
).cuda()
diffusion = GaussianDiffusion(
diffusion_type,
model,
image_size=args.image_size, # Used to be 32
channels=image_channels,
device_of_kernel='cuda',
timesteps=args.time_steps,
loss_type=args.loss_type, #$'l1',
kernel_std=args.kernel_std,
fade_routine=args.fade_routine,
sampling_routine=args.sampling_routine,
discrete=args.discrete
).cuda()
trainer = Trainer(
diffusion,
img_path,
image_size = 32,
train_batch_size = 32,
train_lr = 2e-5,
train_num_steps = 700000, # total training steps
gradient_accumulate_every = 2, # gradient accumulation steps
ema_decay = 0.995, # exponential moving average decay
fp16 = False, # turn on mixed precision training with apex
results_folder = args.save_folder,
load_path = args.load_path
)
if args.test_type == 'train_data':
trainer.test_from_data('train', s_times=args.sample_steps)
elif args.test_type == 'test_data':
trainer.test_from_data('test', s_times=args.sample_steps)
elif args.test_type == 'mixup_train_data':
trainer.test_with_mixup('train')
elif args.test_type == 'mixup_test_data':
trainer.test_with_mixup('test')
elif args.test_type == 'test_random':
trainer.test_from_random('random')
elif args.test_type == 'test_fid_distance_decrease_from_manifold':
trainer.fid_distance_decrease_from_manifold(calculate_fid_given_samples, start=0, end=None)
elif args.test_type == 'test_paper_invert_section_images':
trainer.paper_invert_section_images()
elif args.test_type == 'test_paper_showing_diffusion_images_diff':
trainer.paper_showing_diffusion_images()
|