-
Notifications
You must be signed in to change notification settings - Fork 12
/
utils.py
executable file
·136 lines (119 loc) · 7.89 KB
/
utils.py
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
#June 9, Zheng Xu, [email protected]
# some general functions for all training/testing
# -*- coding: utf-8 -*-
import torch as th
from torch.autograd import Variable
import torchvision as thv
import torchvision.transforms as transforms
import torch.nn as nn
import torch.nn.functional as func
import torch.optim as optim
import torch.backends.cudnn as cudnn
import matplotlib.pyplot as plt
import numpy as np
import argparse
import os
import time
from datetime import datetime
def get_autoencoder_args():
parser = argparse.ArgumentParser(description='PyTorch autoencoder')
parser.add_argument('--use-cpu', action='store_true', help='use CPU')
#training
parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
parser.add_argument('--num-workers', default=4, type=int, help='number of workers to load image')
parser.add_argument('--lr-freq', default=30, type=int, help='learning rate scheduler, 0.1 every x epochs')
parser.add_argument('--momentum', default=0.9, type=float, help='momentum')
parser.add_argument('--weight-decay', default=1e-4, type=float, help='weight decay ')
parser.add_argument('--epochs', default=80, type=int, help='number of total epochs') #80 for imagenet
parser.add_argument('--save-model', default='models/', help='folder to save model')
parser.add_argument('--model', default='vgg19', help='net arch:vgg11,13,16,19')
parser.add_argument('--dataset', default='disentangle', help='dataset: cifar10;cifar100;imagenet32')
parser.add_argument('--content-data', default='data/behance_images/crowd_labels_content', help='data path for content images')
parser.add_argument('--style-data', default='data/behance_images/crowd_labels_media', help='data path for style images')
parser.add_argument('--save-freq', default=1e5, type=int, help='save every x epochs')
parser.add_argument('--print-freq', default=1e5, type=int, help='print every x minibatches')
parser.add_argument('--batch-size', default=16, type=int, help='train minibatch size')
parser.add_argument('--test-batch-size', default=8, type=int, help='test minibatch size')
parser.add_argument('--seed', default=2017, type=int, help='random seed')
parser.add_argument('--test-flag', action='store_true', help='testing')
parser.add_argument('--load-model', default=None, type=str, help='load model for fine tuning or testing')
parser.add_argument('--gpuid', default='0,1,2,3', type=str, help='the GPU IDs')
parser.add_argument('--optm', default='adam', help='discriminator optimizer: padam | sgd | adam | rmsp')
parser.add_argument('--g-optm', default='adam', help='generator optimizer: sgd | adam | rmsp')
parser.add_argument('--adam-b1', default=0.9, type=float, help='adam parameter')
parser.add_argument('--adam-b2', default=0.999, type=float, help='adam parameter')
parser.add_argument('--dropout', default=0, type=float, help='dropout rate')
#parser.add_argument('--logfile', default=None, type=str, help='the log file for showing loss curve')
#test
parser.add_argument('--save-image',default='debug', help='path for saving the styled image')
parser.add_argument('--style-image', default="data/style_trans_images/style/picasso.jpg", help='style image path')
parser.add_argument('--content-image', default="data/style_trans_images/content/house.jpg", help='content image path')
parser.add_argument('--imsize', default=256, type=int, help='image size for input & output')
parser.add_argument('--ae-flag', default='wct', type=str, help='autoencoder architecture')
parser.add_argument('--ae-dep', default='E5-E4', type=str, help='autoencoder enc architecture')
parser.add_argument('--ae-mix', default='mul', type=str, help='mixture flag for disentangled features:mul, add')
parser.add_argument('--enc-model', default='models/wct/vgg_normalised_conv3_1.t7', type=str, help='encoder saved model')
parser.add_argument('--dec-model', default='models/wct/feature_invertor_conv3_1.t7', type=str, help='decoder saved model')
parser.add_argument('--dise-model', default='models/dae_020604_disentangle_2017_wct_sgd0.100_0.0001_dp0.00_10_mb16_8', type=str, help='disentangle layers')
parser.add_argument('--st-layers', default='4', type=str, help='#channels for layers of style encoding')
parser.add_argument('--cnt-layers', default='4', type=str, help='#channels for layers of content encoding')
#GAN, discriminator
parser.add_argument('--d-lr', default=0.1, type=float, help='discriminator learning rate')
#parser.add_argument('--rec-w', default=1, type=float, help='weight for recosntruction loss')
parser.add_argument('--cla-w', default=1, type=float, help='weight for discriminator classification loss')
parser.add_argument('--gan-w', default=1, type=float, help='weight for disciminator adversarial loss')
parser.add_argument('--per-w', default=1, type=float, help='weight for perceptron loss')
parser.add_argument('--gram-w', default=1, type=float, help='weight for gram loss')
parser.add_argument('--cycle-w', default=1, type=float, help='weight for cycle consistency')
#parser.add_argument('--w21', default=1, type=float, help='weight for image 21')
parser.add_argument('--disc-model', default=None, type=str, help='saved discriminator for GAN training')
parser.add_argument('--save-run',default='debug', help='path for saving the intemediate results')
parser.add_argument('--diag-flag',default=None, type=str, help='flag for testing mode, only support batch for release')
parser.add_argument('--gan-ratio', default=0.3, type=float, help='the probability of mixing reconstructin as real image, decreasing with epochs')
parser.add_argument('--gr-freq', default=1, type=float, help='gan ratio decreasing frequency')
parser.add_argument('--train-dec', action='store_true', help='train the decoder')
parser.add_argument('--disc-dep', default=3, type=float, help='depth of discriminator')
#parser.add_argument('--patch-dep', default=3, type=float, help='depth of discriminator')
parser.add_argument('--use-mse', action='store_true', help='use mse for perceptrual & gram loss; otherwise use L1')
parser.add_argument('--use-proj', action='store_true', help='use projection discriminator')
#parser.add_argument('--use-cdisc', action='store_true', help='content info for discriminator')
parser.add_argument('--blur-perc', action='store_true', help='average pooling for the perceptual loss')
parser.add_argument('--dec-last', default='sigmoid', type=str, help='the last layer of decoder: sigmoid | tanh | hard')
parser.add_argument('--pred-gm', default=1, type=float, help='the parameter for predictive step of GAN training, prdict discriminator, use 1 as in original prediction method')
parser.add_argument('--trans-flag', default='adin', type=str, help='the transform method: wct | adin')
parser.add_argument('--base-mode', default='c3', type=str, help='the mode for the base of encoder/decoder: 4;c3;c4')
parser.add_argument('--test-dp', action='store_true', help='dropout while testing')
#parse args
args = parser.parse_args()
return args
def get_model_parameters(net):
n = 0.0
for para in net.parameters():
n += para.numel()
print 'network parameters', n/1.0e6, 'M'
def get_dise_cfg(s): #input is the string
if s == '3':
print 'three dise layers'
return '128,128,128'
elif s == '4':
print 'four dise layers'
return '256,D,256,256,U,128'
elif s == '4w':
print 'four wide dise layers'
return '256,D,512,512,U,256'
elif s == '5':
print 'four dise layers'
return '256,D,256,256,256,U,128'
else:
print 'custom dise layers, returned'
return s
def get_base_dep(s):
if s == 'c3':
return {0,1,2,3}
elif s == 'c4':
return {0,1,2,3,4}
elif s == '4':
return {4}
else:
print 'unknown base_dep mode: use c3 as default'
return {0,1,2,3}