FilesExpand file tree

 master

/

main.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

77 lines (68 loc) · 2.6 KB

 master

/

main.py

Copy path

Top

File metadata and controls

• Code
• Blame

77 lines (68 loc) · 2.6 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

from __future__ import print_function

import torch.backends.cudnn as cudnn

import torch

import torchvision.transforms as transforms

import argparse

import os

import random

import sys

import pprint

import datetime

import dateutil

import dateutil.tz

dir_path = (os.path.abspath(os.path.join(os.path.realpath(__file__), './.')))

sys.path.append(dir_path)

from miscc.datasets import TextDataset

from miscc.config import cfg, cfg_from_file

from miscc.utils import mkdir_p

from trainer import GANTrainer

def parse_args():

parser = argparse.ArgumentParser(description='Train a GAN network')

parser.add_argument('--cfg', dest='cfg_file',

help='optional config file',

default='birds_stage1.yml', type=str)

parser.add_argument('--gpu', dest='gpu_id', type=str, default='0')

parser.add_argument('--data_dir', dest='data_dir', type=str, default='')

parser.add_argument('--manualSeed', type=int, help='manual seed')

args = parser.parse_args()

return args

if __name__ == "__main__":

args = parse_args()

if args.cfg_file is not None:

cfg_from_file(args.cfg_file)

if args.gpu_id != -1:

cfg.GPU_ID = args.gpu_id

if args.data_dir != '':

cfg.DATA_DIR = args.data_dir

print('Using config:')

pprint.pprint(cfg)

if args.manualSeed is None:

args.manualSeed = random.randint(1, 10000)

random.seed(args.manualSeed)

torch.manual_seed(args.manualSeed)

if cfg.CUDA:

torch.cuda.manual_seed_all(args.manualSeed)

now = datetime.datetime.now(dateutil.tz.tzlocal())

timestamp = now.strftime('%Y_%m_%d_%H_%M_%S')

output_dir = '../output/%s_%s_%s' % \

(cfg.DATASET_NAME, cfg.CONFIG_NAME, timestamp)

num_gpu = len(cfg.GPU_ID.split(','))

if cfg.TRAIN.FLAG:

image_transform = transforms.Compose([

transforms.RandomCrop(cfg.IMSIZE),

transforms.RandomHorizontalFlip(),

transforms.ToTensor(),

transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

dataset = TextDataset(cfg.DATA_DIR, 'train',

imsize=cfg.IMSIZE,

transform=image_transform)

assert dataset

dataloader = torch.utils.data.DataLoader(

dataset, batch_size=cfg.TRAIN.BATCH_SIZE * num_gpu,

drop_last=True, shuffle=True, num_workers=int(cfg.WORKERS))

algo = GANTrainer(output_dir)

algo.train(dataloader, cfg.STAGE)

else:

datapath= '%s/test/val_captions.t7' % (cfg.DATA_DIR)

algo = GANTrainer(output_dir)

algo.sample(datapath, cfg.STAGE)