File size: 1,822 Bytes
da855ff |
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 |
# This code is based on https://github.com/openai/guided-diffusion
"""
Train a diffusion model on images.
"""
import os
import json
from utils.fixseed import fixseed
from utils.parser_util import train_args
from utils import dist_util
from train.training_loop import TrainLoop
from data_loaders.get_data import get_dataset_loader
from utils.model_util import create_model_and_diffusion
from train.train_platforms import ClearmlPlatform, TensorboardPlatform, NoPlatform # required for the eval operation
def main():
args = train_args()
fixseed(args.seed)
train_platform_type = eval(args.train_platform_type)
train_platform = train_platform_type(args.save_dir)
train_platform.report_args(args, name='Args')
if args.save_dir is None:
raise FileNotFoundError('save_dir was not specified.')
elif os.path.exists(args.save_dir) and not args.overwrite:
raise FileExistsError('save_dir [{}] already exists.'.format(args.save_dir))
elif not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
args_path = os.path.join(args.save_dir, 'args.json')
with open(args_path, 'w') as fw:
json.dump(vars(args), fw, indent=4, sort_keys=True)
dist_util.setup_dist(args.device)
print("creating data loader...")
data = get_dataset_loader(name=args.dataset, batch_size=args.batch_size, num_frames=args.num_frames)
print("creating model and diffusion...")
model, diffusion = create_model_and_diffusion(args, data)
model.to(dist_util.dev())
model.rot2xyz.smpl_model.eval()
print('Total params: %.2fM' % (sum(p.numel() for p in model.parameters_wo_clip()) / 1000000.0))
print("Training...")
TrainLoop(args, train_platform, model, diffusion, data).run_loop()
train_platform.close()
if __name__ == "__main__":
main()
|