def train():
X = T.tensor4('images')
y = T.lvector('labels')
X_ = nn.placeholder((bs, ) + image_shape, name='images_plhd')
y_ = nn.placeholder((bs, ), dtype='int64', name='labels_plhd')
net = nets[architecture]((None, ) + image_shape)
nn.set_training_on()
updates, losses, grad_norms = net.get_updates(X, y, l2_coeff, lr)
train_net = nn.function([],
list(losses.values()),
updates=updates,
givens={
X: X_,
y: y_
},
name='train net')
nn.set_training_off()
err, losses = net.get_accuracy(X, y)
valid_net = nn.function([], [err, losses['loss']],
givens={
X: X_,
y: y_
},
name='validate net')
train_data = CIFAR10((X_train, y_train), (X_, y_),
bs,
n_epochs,
training=True,
shuffle=True,
augmentation=transforms)
valid_data = CIFAR10((X_test, y_test), (X_, y_),
bs,
1,
training=False,
shuffle=False)
mon = nn.Monitor(model_name=architecture, print_freq=print_freq)
for it in train_data:
with mon:
losses_ = train_net()
if np.any(np.isnan(losses_)) or np.any(np.isinf(losses_)):
raise ValueError('NAN loss!')
for j, k in enumerate(losses.keys()):
mon.plot(k, losses_[j])
if it % valid_freq == 0:
mean_res = np.mean([valid_net() for _ in valid_data], 0)
mon.plot('validation error', mean_res[0])
mon.plot('validation loss', mean_res[1])
mon.dump(nn.utils.shared2numpy(net.params),
'%s.npy' % architecture, 5)
print('Training finished!')
def train(input_shape,
output_size,
image_path,
bs=64,
z_dim=1024,
n_iters=int(1e6)):
batch_input_shape = (bs, ) + input_shape
net = LSGAN((bs, z_dim), output_size)
X__ = T.tensor4('image', 'float32')
X = DataManager.preprocess(X__)
z = srng.uniform((bs, z_dim), -1, 1, ndim=2, dtype='float32')
X_ = theano.shared(np.zeros(batch_input_shape, 'float32'),
'image placeholder')
nn.set_training_status(True)
gen_loss, dis_loss, updates_gen, updates_dis = net.learn(X, z, 1e-3, .5)
train_gen = nn.function([],
gen_loss,
updates=updates_gen,
name='train generator')
train_dis = nn.function([],
dis_loss,
updates=updates_dis,
givens={X__: X_},
name='train discriminator')
nn.set_training_status(False)
fixed_noise = T.constant(np.random.uniform(-1, 1, (bs, z_dim)),
'fixed noise', 2, 'float32')
gen_imgs = net(fixed_noise)
generate = nn.function([], gen_imgs, name='generate images')
n_epochs = int(bs * n_iters / len(os.listdir(image_path)))
dm = DataManager(output_size, X_, image_path, bs, n_epochs, True)
mon = nn.monitor.Monitor(model_name='LSGAN', use_visdom=True)
batches = dm.get_batches()
start = time.time()
for it in batches:
dis_loss_ = train_dis()
gen_loss_ = train_gen()
if np.isnan(gen_loss_ + dis_loss_) or np.isinf(dis_loss_ + gen_loss_):
raise ValueError('Training failed! Stopped.')
mon.plot('discriminator loss', dis_loss_)
mon.plot('generator loss', gen_loss_)
if it % 500 == 499:
fake_imgs = generate()
mon.imwrite('fake images', fake_imgs / 2. + .5)
mon.plot('time elapsed', (time.time() - start) / 60.)
nn.utils.save(net, mon.current_folder + '/training.pkl')
mon.flush()
mon.tick()
net.save_params(mon.current_folder + '/params.npz')
print('Training finished')
def train():
enc = VGG19(input_shape)
decs = [Decoder(enc, i, name='decoder %d' % i) for i in indices]
sty_net = StyleTransfer(enc, decs)
X = T.tensor4('input')
Y = T.tensor4('style')
idx = T.scalar('iter', 'int32')
X_ = nn.placeholder((bs,) + input_shape[1:], name='input_plhd')
Y_ = nn.placeholder((bs,) + input_shape[1:], name='style_plhd')
lr_ = nn.placeholder(value=lr, name='lr_plhd')
nn.set_training_on()
losses = [dec.cost(X) for dec in decs]
updates = [nn.adam(loss[0] + weight * loss[1], dec.trainable, lr) for loss, dec in zip(losses, decs)]
nn.anneal_learning_rate(lr_, idx, 'inverse', decay=decay)
trains = [nn.function([], [loss[0], loss[1], dec(X, True)], givens={X: X_}, updates=update, name='train decoder')
for loss, dec, update in zip(losses, decs, updates)]
nn.set_training_off()
X_styled = sty_net(X, Y)
transfer = nn.function([], X_styled, givens={X: X_, Y: Y_}, name='transfer style')
data_train = DataManager(X_, input_path_train, bs, n_epochs, True, num_val_imgs=num_val_imgs, input_shape=input_shape)
data_test = DataManagerStyleTransfer((X_, Y_), (input_path_val, style_path_val), bs, 1, input_shape=input_shape)
mon = nn.Monitor(model_name='WCT', valid_freq=print_freq)
print('Training...')
for it in data_train:
results = [train(it) for train in trains]
with mon:
for layer, res in zip(indices, results):
if np.isnan(res[0] + res[1]) or np.isinf(res[0] + res[1]):
raise ValueError('Training failed!')
mon.plot('pixel loss at layer %d' % layer, res[0])
mon.plot('feature loss at layer %d' % layer, res[1])
if it % val_freq == 0:
mon.imwrite('recon img at layer %d' % layer, res[2])
for i in data_test:
img = transfer()
mon.imwrite('stylized image %d' % i, img)
mon.imwrite('input %d' % i, X_.get_value())
mon.imwrite('style %d' % i, Y_.get_value())
for idx, dec in zip(indices, decs):
mon.dump(nn.utils.shared2numpy(dec.params), 'decoder-%d.npz' % idx, 5)
mon.flush()
for idx, dec in zip(indices, decs):
mon.dump(nn.utils.shared2numpy(dec.params), 'decoder-%d-final.npz' % idx)
print('Training finished!')
def test_tracking():
trivial_loops = 5
shape = (3, 3)
a = T.matrix('a')
b = T.scalar('b')
b_ = nn.placeholder(value=1.)
def foo(x, y):
for i in range(trivial_loops):
nn.track('loop %d' % (i + 1), x)
x += y
return x
c = foo(a, b)
func = nn.function([a], c, givens={b: b_})
a_ = np.zeros(shape, 'float32')
res_numpy = [a_ + i for i in range(trivial_loops + 1)]
utt.assert_allclose(func(a_), res_numpy[-1])
for i in range(trivial_loops):
trackeds = nn.eval_tracked_vars({a: a_, b: b_.get_value()})
assert all(
np.allclose(x, y)
for x, y in zip(list(trackeds.values()), res_numpy[:-1]))
def test_lr_annealing():
base_lr = 1.
n_iters = 50
decay = 1e-2
step = 10
def anneal_learning_rate(lr, t, method='half-life', **kwargs):
lr_ = lr
if method == 'half-life':
num_iters = kwargs.pop('num_iters', None)
decay = kwargs.pop('decay', .1)
if num_iters is None:
raise ValueError('num_iters must be provided.')
if (t == num_iters // 2) or (t == 3 * num_iters // 4):
return lr_ * decay
else:
return lr_
elif method == 'step':
step = kwargs.pop('step', None)
decay = kwargs.pop('decay', .5)
if step is None:
raise ValueError('step must be provided.')
if t % step == 0:
return lr_ * decay
else:
return lr_
elif method == 'exponential':
decay = kwargs.pop('decay', 1e-4)
return lr_ * np.exp(-decay * t)
elif method == 'linear':
num_iters = kwargs.pop('num_iters', None)
if num_iters is None:
raise ValueError('num_iters must be provided.')
return lr_ * (1. - t / np.cast[theano.config.floatX](num_iters))
elif method == 'inverse':
decay = kwargs.pop('decay', .01)
return lr_ / (1. + decay * t)
else:
raise ValueError('Unknown annealing method.')
idx = T.scalar('it', 'int32')
for method in ('linear', 'step', 'exponential', 'inverse', 'half-life'):
print('Testing method %s' % method)
lr_ = nn.placeholder(dtype=theano.config.floatX, value=base_lr, name='lr')
y = 0. + lr_
nn.anneal_learning_rate(lr_, idx, method, num_iters=n_iters, decay=decay, step=step)
func = nn.function([idx], y)
vals_th, vals_np = [], []
lr = base_lr
for it in range(n_iters):
func(it + 1)
vals_th.append(lr_.get_value())
lr = anneal_learning_rate(lr if method in ('step', 'half-life') else base_lr, it+1, method,
num_iters=n_iters, decay=decay, step=step)
vals_np.append(lr)
assert_allclose(vals_th, vals_np)
def style_transfer():
enc = VGG19(input_shape)
decs = [Decoder(enc, i, name='decoder %d' % i) for i in indices]
sty_net = StyleTransfer(enc, decs)
X = T.tensor4('input')
Y = T.tensor4('style')
mon = nn.Monitor(current_folder=args.path_to_weight_files)
for idx, layer in enumerate(indices):
weights = mon.load('decoder-%d-final.npz' % layer)
nn.utils.numpy2shared(weights, decs[idx].params)
nn.set_training_off()
X_styled = sty_net(X, Y)
transfer = nn.function([X, Y], X_styled, name='transfer style')
if os.path.isfile(input_path_test) and os.path.isfile(style_path_test):
input = prep_image_test(misc.imread(input_path_test))
style = prep_image_test(misc.imread(style_path_test))
output = transfer(input, style)
mon.imwrite('test %s' % input_path_test[:-4], input)
mon.imwrite('test %s' % style_path_test[:-4], style)
mon.imwrite('test %s %s' % (input_path_test[:-4], style_path_test[:-4]), output)
elif os.path.isfile(input_path_test) and os.path.isdir(style_path_test):
input = prep_image_test(misc.imread(input_path_test))
style_files = os.listdir(style_path_test)
for style_file in style_files:
style = prep_image_test(misc.imread(os.path.join(style_path_test, style_file)))
output = transfer(input, style)
mon.imwrite('test %s' % style_file[:-4], style)
mon.imwrite('test %s %s' % (input_path_test[:-4], style_file[:-4]), output)
mon.imwrite('test %s' % input_path_test[:-4], input)
elif os.path.isdir(input_path_test) and os.path.isfile(style_path_test):
style = prep_image_test(misc.imread(style_path_test))
input_files = os.listdir(input_path_test)
for input_file in input_files:
input = prep_image_test(misc.imread(os.path.join(input_path_test, input_file)))
output = transfer(input, style)
mon.imwrite('test %s' % input_file[:-4], input)
mon.imwrite('test %s %s' % (input_file[:-4], style_path_test[:-4]), output)
mon.imwrite('test %s' % style_path_test[:-4], style)
else:
style_files = os.listdir(style_path_test)
input_files = os.listdir(input_path_test)
for style_file in style_files:
style = prep_image_test(misc.imread(os.path.join(style_path_test, style_file)))
for input_file in input_files:
input = prep_image_test(misc.imread(os.path.join(input_path_test, input_file)))
output = transfer(input, style)
mon.imwrite('test %s' % input_file[:-4], input)
mon.imwrite('test %s %s' % (input_file[:-4], style_file[:-4]), output)
mon.imwrite('test %s' % style_file[:-4], style)
mon.flush()
print('Testing finished!')
def test_yiq():
from scipy import misc
im = misc.imread('test_files/lena.jpg').astype(theano.config.floatX) / 255.
im = np.transpose(im[None], (0, 3, 1, 2))
a = T.tensor4('input')
b = nn.utils.rgb2yiq(a)
c = nn.utils.yiq2rgb(b)
func = nn.function([a], [b[:, 0:1], c])
im_out_y, im_out = func(im)
misc.imsave('test_files/lena_from_yiq.jpg', np.transpose(im_out[0], (1, 2, 0)))
misc.imsave('test_files/lena_y.jpg', im_out_y[0, 0])
def test_diff_of_gaussians():
size1 = 21
size2 = 51
x = T.tensor4('image')
x1 = nn.utils.difference_of_gaussian(x, size1)
x2 = nn.utils.difference_of_gaussian(x, size2)
func = nn.function([x], [x1, x2])
from scipy import misc
im = misc.imread('test_files/lena_small.png').astype(theano.config.floatX) / 255.
im = np.transpose(im[None], (0, 3, 1, 2))
x1_, x2_ = func(im)
misc.imsave('test_files/lena_small_dog1.jpg', np.transpose(x1_[0], (1, 2, 0)))
misc.imsave('test_files/lena_small_dog2.jpg', np.transpose(x2_[0], (1, 2, 0)))
def test_downsample():
shape = (1, 3, 220, 220)
kernel_type = 'gauss1sq2'
x = T.tensor4('input')
downsample = nn.DownsamplingLayer(shape, 4, kernel_type=kernel_type)
y = downsample(x)
f = nn.function([x], y)
from imageio import imwrite, imread
x = imread('test_files/lena_small.png').astype('float32') / 255.
x = np.transpose(x[None], (0, 3, 1, 2))
out = f(x)
out = np.transpose(out[0], (1, 2, 0))
imwrite('test_files/lena_small_downsampled.jpg', out)
def test_partial_conv_based_padding():
shape = (1, 3, 5, 5)
num_filters = 2
filter_size = 3
x = T.ones(shape)
conv = nn.Conv2DLayer(shape,
num_filters,
filter_size,
border_mode='partial')
y = conv(x)
f = nn.function([], y)
print(f())
print(conv.update_mask.eval())
print(conv.mask_ratio)
def test_pearsonr():
size = (64, 1000)
pred = T.matrix('pred')
gt = T.matrix('gt')
corr = nn.pearsonr(pred, gt)
f = nn.function([pred, gt], corr)
gt = np.random.normal(size=size).astype(theano.config.floatX)
pred = np.random.uniform(size=size).astype(theano.config.floatX)
c = f(pred, gt)
from scipy import stats
c_ref, _ = stats.pearsonr(pred.flatten(), gt.flatten())
utt.assert_allclose(c_ref, c)
def test_upsampling_layer():
input_shape = (1, 3, 220, 220)
X = T.tensor4('input')
bilinear_up = nn.UpsamplingLayer(input_shape, 4, method='bilinear')
nearest_up = nn.UpsamplingLayer(input_shape, 4, method='nearest')
x_bilinear = bilinear_up(X)
x_nearest = nearest_up(X)
func = nn.function([X], [x_bilinear, x_nearest])
from scipy import misc
im = np.transpose(misc.imread('test_files/lena_small.png').astype(theano.config.floatX), (2, 0, 1))[None] / 255.
x_bi_, x_ne_ = func(im)
x_bi_ = np.transpose(x_bi_[0], (1, 2, 0))
x_ne_ = np.transpose(x_ne_[0], (1, 2, 0))
misc.imsave('test_files/lena_small_bilinear.jpg', x_bi_)
misc.imsave('test_files/lena_small_nearest.jpg', x_ne_)
def test_meshgrid():
height = 10.
width = 100.
h = T.scalar()
w = T.scalar()
x, y = nn.utils.linspace(0, w, w), nn.utils.linspace(0, h, h)
X, Y = nn.utils.meshgrid(x, y)
f = nn.function([w, h], [X, Y])
X_th, Y_th = f(width, height)
x, y = np.linspace(0, width, width), np.linspace(0, height, height)
X_np, Y_np = np.meshgrid(x, y)
utt.assert_allclose(X_th, X_np)
utt.assert_allclose(Y_th, Y_np)
def test_rfft2():
X = T.tensor4('input')
Y = nn.utils.rfft2(X)
Y_i = nn.utils.irfft2(Y)
rfft2 = nn.function([X], [Y, Y_i])
from imageio import imread
x = imread('test_files/lena_small.png').astype('float32') / 255.
x = np.transpose(x, (2, 0, 1))
x_fft_np = np.fft.rfft2(x)
x = x[None]
x_fft_th, x_ifft_th = rfft2(x)
x_fft_th = x_fft_th[0]
utt.assert_allclose(np.real(x_fft_np), x_fft_th[..., 0])
utt.assert_allclose(np.imag(x_fft_np), x_fft_th[..., 1])
np.allclose(x_ifft_th, x, 1e-4)
def test_frac_bilinear_upsampling():
frac_ratio1 = ((5, 2), (11, 5))
frac_ratio2 = ((2, 3), (5, 13))
frac_ratio3 = ((2, 3), (17, 7))
X = T.tensor4('input')
Y1 = nn.utils.frac_bilinear_upsampling(X, frac_ratio1)
Y2 = nn.utils.frac_bilinear_upsampling(X, frac_ratio2)
Y3 = nn.utils.frac_bilinear_upsampling(X, frac_ratio3)
f = nn.function([X], [Y1, Y2, Y3])
from imageio import imread, imwrite
x = imread('test_files/lena_small.png').astype('float32') / 255.
x = np.transpose(x, (2, 0, 1))[None]
y1, y2, y3 = f(x)
out = np.array(np.transpose(y1[0], (1, 2, 0)))
imwrite('test_files/lena_small_frac_bi_up.jpg', out)
out = np.array(np.transpose(y2[0], (1, 2, 0)))
imwrite('test_files/lena_small_frac_bi_down.jpg', out)
out = np.array(np.transpose(y3[0], (1, 2, 0)))
imwrite('test_files/lena_small_frac_bi_mixed.jpg', out)
input_x = np.array(
[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[9, 10], [11, 12]]],
ndmin=4).astype(theano.config.floatX)
up_x = nn.utils.frac_bilinear_upsampling(input=input_x,
frac_ratio=((7, 4), (5, 3)))
num_up_x = np.array([[[[1., 1.2, 1.8, 2.],
[1.28571429, 1.48571429, 2.08571429, 2.28571429],
[2.42857143, 2.62857143, 3.22857143, 3.42857143],
[3., 3.2, 3.8, 4.]],
[[5., 5.2, 5.8, 6.],
[5.28571429, 5.48571429, 6.08571429, 6.28571429],
[6.42857143, 6.62857143, 7.22857143, 7.42857143],
[7., 7.2, 7.8, 8.]],
[[9., 9.2, 9.8, 10.],
[9.28571429, 9.48571429, 10.08571429, 10.28571429],
[
10.42857143, 10.62857143, 11.22857143,
11.42857143
], [11., 11.2, 11.8,
12.]]]]).astype(theano.config.floatX)
utt.assert_allclose(up_x.eval(), num_up_x, rtol=1e-6)
def test_model_zoo_vgg19():
top = 1
root = 'test_files/'
weight_file = root + 'vgg19_from_pytorch.npz'
imagenet_classes_file = root + 'imagenet_classes.txt'
image_list = [
'ILSVRC2012_val_00000001.JPEG', 'ILSVRC2012_val_00000002.JPEG',
'ILSVRC2012_val_00000003.JPEG', 'ILSVRC2012_val_00000004.JPEG',
'ILSVRC2012_val_00000005.JPEG'
]
net = nn.model_zoo.VGG19((None, 3, 224, 224))
net.load_params(weight_file)
X = T.tensor4('input')
test = nn.function([X], net(X))
with open(imagenet_classes_file, 'r') as f:
classes = [s.strip() for s in f.readlines()]
mean_rgb = np.array([0.485, 0.456, 0.406],
dtype=theano.config.floatX)[None, None, :]
std_rgb = np.array([0.229, 0.224, 0.225],
dtype=theano.config.floatX)[None, None, :]
results = []
for fname in image_list:
print(fname)
rawim, im = nn.utils.prep_image2(root + fname, mean_rgb, std_rgb)
prob = test(im)[0]
res = sorted(zip(classes, prob), key=lambda t: t[1],
reverse=True)[:top]
for c, p in res:
print(' ', c, p)
results.append(c)
print('\n')
sample_res = [
'rock python, rock snake, Python sebae', 'ski',
'Shetland sheepdog, Shetland sheep dog, Shetland', 'soup bowl',
'bassinet'
]
assert results == sample_res
def test_random():
enc = Encoder((None, ) + input_size, vgg_param_file)
dec = Decoder(enc.output_shape, dec_param_file)
X = T.tensor4('input')
Y = T.tensor4('style')
weights = [T.vector('weights') for i in range(len(vgg_info))]
nn.set_training_off()
X_styled = dec(enc((X, Y), weights))
test = nn.function([X, Y] + weights, X_styled, name='test generator')
style_folder = os.listdir(style_img_folder)
input_folder = os.listdir(input_img_folder)
time_list = []
if not os.path.exists('outputs'):
os.mkdir('outputs')
for style_file in style_folder:
sty_img = prep_image_test(
misc.imread(style_img_folder + '/' + style_file))
for input_file in input_folder:
try:
input_img = prep_image_test(
misc.imread(input_img_folder + '/' + input_file))
except ValueError:
continue
for i in range(num_styles):
start = time.time()
output = test(input_img, sty_img, *get_weights(vgg_info))
time_list.append(time.time() - start)
output = np.transpose(output[0], (1, 2, 0))
misc.imsave(
os.path.join(
'outputs', input_file[:-4] + '_' + style_file[:-4] +
'_%d.jpg' % i), output)
print('Took %f s/image' % np.mean(time_list))
print('Testing finished!')
def train_sngan(z_dim=128, image_shape=(3, 32, 32), bs=64, n_iters=int(1e5)):
gen = DCGANGenerator((None, z_dim))
dis = SNDCGANDiscriminator(gen.output_shape)
z = srng.uniform((bs, z_dim), -1, 1, ndim=2, dtype='float32')
X = T.tensor4('image', 'float32')
X_ = theano.shared(np.zeros((bs, ) + image_shape, 'float32'),
'image_placeholder')
# training
nn.set_training_status(True)
X_fake = gen(z)
y_fake = dis(X_fake)
y_real = dis(X)
dis_loss_real = T.mean(T.nnet.softplus(-y_real))
dis_loss_fake = T.mean(T.nnet.softplus(y_fake))
dis_loss = dis_loss_real + dis_loss_fake
gen_loss = T.mean(T.nnet.softplus(-y_fake))
updates_gen = nn.adam(gen_loss, gen.trainable, args.adam_alpha,
args.adam_beta1, args.adam_beta2)
updates_dis = nn.adam(dis_loss, dis.trainable, args.adam_alpha,
args.adam_beta1, args.adam_beta2)
train_gen = nn.function([],
gen_loss,
updates=updates_gen,
name='train generator')
train_dis = nn.function([],
dis_loss,
updates=updates_dis,
givens={X: X_},
name='train discriminator')
# testing
nn.set_training_status(False)
fixed_noise = T.constant(np.random.uniform(-1, 1, (bs, z_dim)),
'fixed noise', 2, 'float32')
gen_imgs = gen(fixed_noise)
generate = nn.function([], gen_imgs, name='generate images')
dm = DataManager(X_, n_iters, bs, True)
mon = nn.monitor.Monitor(model_name='LSGAN', use_visdom=args.use_visdom)
epoch = 0
print('Training...')
batches = dm.get_batches(epoch, dm.n_epochs, infinite=True)
start = time.time()
for iteration in range(n_iters):
#update generator
training_gen_cost = train_gen()
if np.isnan(training_gen_cost) or np.isinf(training_gen_cost):
raise ValueError('Training failed due to NaN cost')
mon.plot('training gen cost', training_gen_cost)
#update discriminator
training_disc_cost = []
for i in range(args.n_dis):
batches.__next__()
training_disc_cost.append(train_dis())
if np.isnan(training_disc_cost[-1]) or np.isinf(
training_disc_cost[-1]):
raise ValueError('Training failed due to NaN cost')
mon.plot('training disc cost', np.mean(training_disc_cost))
if iteration % args.valid_freq == 0:
gen_images = generate()
mon.imwrite('generated image', dm.unnormalize(gen_images))
mon.plot('time elapsed', (time.time() - start) / 60.)
mon.flush()
mon.tick()
mon.flush()
print('Training finished!')
def train():
X_A_full = T.tensor4('A')
X_B_full = T.tensor4('B')
X_A = pre_process(X_A_full)
X_B = pre_process(X_B_full)
z = nn.utils.srng.normal((bs, latent_dim))
idx = T.scalar('iter')
X_A_ = nn.placeholder((bs, 3, image_size*4, image_size*4), name='A_plhd')
X_B_ = nn.placeholder((bs, 3, image_size*4, image_size*4), name='B_plhd')
lr_ = nn.placeholder(value=lr, name='lr_plhd')
net = AugmentedCycleGAN((None, 3, image_size, image_size), latent_dim, n_gen_filters, n_dis_filters, n_enc_filters, 3,
use_dropout, use_sigmoid, use_latent_gan)
nn.set_training_on()
updates_dis, updates_gen, dis_losses, dis_preds, gen_losses, grad_norms = net.learn(X_A, X_B, z, lambda_A, lambda_B,
lambda_z_B, lr=lr_, beta1=beta1,
max_norm=max_norm)
train_dis = nn.function([], list(dis_losses.values()), updates=updates_dis, givens={X_A_full: X_A_, X_B_full: X_B_},
name='train discriminators')
train_gen = nn.function([], list(gen_losses.values()), updates=updates_gen, givens={X_A_full: X_A_, X_B_full: X_B_},
name='train generators')
discriminate = nn.function([], list(dis_preds.values()), givens={X_A_full: X_A_, X_B_full: X_B_}, name='discriminate')
compute_grad_norms = nn.function([], list(grad_norms.values()), givens={X_A_full: X_A_, X_B_full: X_B_},
name='compute grad norms')
nn.anneal_learning_rate(lr_, idx, 'linear', num_iters=n_epochs_decay)
train_dis_decay = nn.function([idx], list(dis_losses.values()), updates=updates_dis, givens={X_A_full: X_A_, X_B_full: X_B_},
name='train discriminators with decay')
nn.set_training_off()
fixed_z = T.constant(np.random.normal(size=(bs, latent_dim)), dtype='float32')
fixed_multi_z = T.constant(np.repeat(np.random.normal(size=(n_multi, latent_dim)), bs, 0), dtype='float32')
visuals = net.generate_cycle(X_A, X_B, fixed_z)
multi_fake_B = net.generate_multi(X_A, fixed_multi_z)
visualize_single = nn.function([], list(visuals.values()), givens={X_A_full: X_A_, X_B_full: X_B_}, name='visualize single')
visualize_multi = nn.function([], multi_fake_B, givens={X_A_full: X_A_}, name='visualize multi')
train_data = Edges2Shoes((X_A_, X_B_), bs, n_epochs + n_epochs_decay + 1, 'train', True)
val_data = Edges2Shoes((X_A_, X_B_), bs, 1, 'val', False, num_data=bs)
mon = nn.Monitor(model_name='Augmented_CycleGAN', print_freq=print_freq)
print('Training...')
for it in train_data:
epoch = 1 + it // (len(train_data) // bs)
with mon:
res_dis = train_dis() if epoch <= n_epochs else train_dis_decay(epoch - n_epochs)
res_gen = train_gen()
preds = discriminate()
grads_ = compute_grad_norms()
mon.plot('lr', lr_.get_value())
for j, k in enumerate(dis_losses.keys()):
mon.plot(k, res_dis[j])
for j, k in enumerate(gen_losses.keys()):
mon.plot(k, res_gen[j])
for j, k in enumerate(dis_preds.keys()):
mon.hist(k, preds[j])
for j, k in enumerate(grad_norms.keys()):
mon.plot(k, grads_[j])
if it % valid_freq == 0:
for _ in val_data:
vis_single = visualize_single()
vis_multi = visualize_multi()
for j, k in enumerate(visuals.keys()):
mon.imwrite(k, vis_single[j][:n_imgs_to_save], callback=unnormalize)
for j, fake_B in enumerate(vis_multi):
mon.imwrite('fake_B_multi_%d.jpg' % j, fake_B, callback=unnormalize)
mon.dump(nn.utils.shared2numpy(net.netG_A_B.params), 'gen_A_B.npy', 5)
mon.dump(nn.utils.shared2numpy(net.netG_B_A.params), 'gen_B_A.npy', 5)
mon.dump(nn.utils.shared2numpy(net.netD_A.params), 'dis_A.npy', 5)
mon.dump(nn.utils.shared2numpy(net.netD_B.params), 'dis_B.npy', 5)
mon.dump(nn.utils.shared2numpy(net.netE_B.params), 'enc_B.npy', 5)
if use_latent_gan:
mon.dump(nn.utils.shared2numpy(net.netD_z_B.params), 'dis_z_B.npy', 5)
mon.flush()
mon.dump(nn.utils.shared2numpy(net.netG_A_B.params), 'gen_A_B.npy')
mon.dump(nn.utils.shared2numpy(net.netG_B_A.params), 'gen_B_A.npy')
mon.dump(nn.utils.shared2numpy(net.netD_A.params), 'dis_A.npy')
mon.dump(nn.utils.shared2numpy(net.netD_B.params), 'dis_B.npy')
mon.dump(nn.utils.shared2numpy(net.netE_B.params), 'enc_B.npy')
if use_latent_gan:
mon.dump(nn.utils.shared2numpy(net.netD_z_B.params), 'dis_z_B.npy')
print('Training finished!')
if __name__ == '__main__':
from theano.tensor.nnet import conv2d
from theano.tensor.nnet.conv3d2d import conv3d
x = T.tensor4()
h = T.tensor4()
# conv2d_temp = []
# for patch_num in range(36):
# input = T.concatenate([x[patch_num, :, :, :].dimshuffle(0, 'x', 1, 2)]*512, 1)
# filter = h[patch_num, :, :, :].dimshuffle('x', 0, 1, 2)
# conv2d_temp.append(conv2d(input, filter,
# border_mode='half'
# ))
# conv_out = T.concatenate(conv2d_temp, axis=0)
x_con = T.concatenate([x] * 512, 1)
x_prime = x_con.dimshuffle(0, 'x', 1, 2, 3)
conv_out3 = conv3d(x_prime, h.dimshuffle('x', 0, 1, 2, 3), border_mode='half')
conv_out3 = T.stack([conv_out3[i, i] for i in range(36)])
# func = nn.function([x, h], conv_out)
func_3 = nn.function([x, h], conv_out3)
x = np.random.rand(36, 1, 10, 10).astype('float32')
h = np.random.rand(36, 512, 3, 3).astype('float32')
print(func_3(x, h).shape)
# print(func(x, h).shape)
# assert np.allclose(func(x, h), func_3(x, h))
