def main():
"""
Main script.
Steps:
* Get and set context.
* Load Dataset
* Initialize DataIterator.
* Create Networks
* Net for Labeled Data
* Net for Unlabeled Data
* Net for Test Data
* Create Solver.
* Training Loop.
* Test
* Training
* by Labeled Data
* Calculate Cross Entropy Loss
* by Unlabeled Data
* Estimate Adversarial Direction
* Calculate LDS Loss
"""
args = get_args()
# Get context.
from nnabla.contrib.context import extension_context
extension_module = args.context
if args.context is None:
extension_module = 'cpu'
logger.info("Running in %s" % extension_module)
ctx = extension_context(extension_module, device_id=args.device_id)
nn.set_default_context(ctx)
shape_x = (1, 28, 28)
n_h = args.n_units
n_y = args.n_class
# Load MNist Dataset
from mnist_data import MnistDataSource
with MnistDataSource(train=True) as d:
x_t = d.images
t_t = d.labels
with MnistDataSource(train=False) as d:
x_v = d.images
t_v = d.labels
x_t = np.array(x_t / 256.0).astype(np.float32)
x_t, t_t = x_t[:args.n_train], t_t[:args.n_train]
x_v, t_v = x_v[:args.n_valid], t_v[:args.n_valid]
# Create Semi-supervised Datasets
x_l, t_l, x_u, _ = split_dataset(x_t, t_t, args.n_labeled, args.n_class)
x_u = np.r_[x_l, x_u]
x_v = np.array(x_v / 256.0).astype(np.float32)
# Create DataIterators for datasets of labeled, unlabeled and validation
di_l = DataIterator(args.batchsize_l, [x_l, t_l])
di_u = DataIterator(args.batchsize_u, [x_u])
di_v = DataIterator(args.batchsize_v, [x_v, t_v])
# Create networks
# feed-forward-net building function
def forward(x, test=False):
return mlp_net(x, n_h, n_y, test)
# Net for learning labeled data
xl = nn.Variable((args.batchsize_l,) + shape_x, need_grad=False)
hl = forward(xl, test=False)
tl = nn.Variable((args.batchsize_l, 1), need_grad=False)
loss_l = F.mean(F.softmax_cross_entropy(hl, tl))
# Net for learning unlabeled data
xu = nn.Variable((args.batchsize_u,) + shape_x, need_grad=False)
r = nn.Variable((args.batchsize_u,) + shape_x, need_grad=True)
eps = nn.Variable((args.batchsize_u,) + shape_x, need_grad=False)
loss_u, yu = vat(xu, r, eps, forward, distance)
# Net for evaluating valiation data
xv = nn.Variable((args.batchsize_v,) + shape_x, need_grad=False)
hv = forward(xv, test=True)
tv = nn.Variable((args.batchsize_v, 1), need_grad=False)
# Create solver
solver = S.Adam(args.learning_rate)
solver.set_parameters(nn.get_parameters())
# Monitor trainig and validation stats.
import nnabla.monitor as M
monitor = M.Monitor(args.model_save_path)
monitor_verr = M.MonitorSeries("Test error", monitor, interval=240)
monitor_time = M.MonitorTimeElapsed("Elapsed time", monitor, interval=240)
# Training Loop.
t0 = time.time()
for i in range(args.max_iter):
# Validation Test
if i % args.val_interval == 0:
n_error = calc_validation_error(
di_v, xv, tv, hv, args.val_iter)
monitor_verr.add(i, n_error)
#################################
## Training by Labeled Data #####
#################################
# input minibatch of labeled data into variables
xl.d, tl.d = di_l.next()
# initialize gradients
solver.zero_grad()
# forward, backward and update
loss_l.forward(clear_no_need_grad=True)
loss_l.backward(clear_buffer=True)
solver.weight_decay(args.weight_decay)
solver.update()
#################################
## Training by Unlabeled Data ###
#################################
# input minibatch of unlabeled data into variables
xu.d, = di_u.next()
##### Calculate Adversarial Noise #####
# Sample random noise
n = np.random.normal(size=xu.shape).astype(np.float32)
# Normalize noise vector and input to variable
r.d = get_direction(n)
# Set xi, the power-method scaling parameter.
eps.data.fill(args.xi_for_vat)
# Calculate y without noise, only once.
yu.forward(clear_buffer=True)
# Do power method iteration
for k in range(args.n_iter_for_power_method):
# Initialize gradient to receive value
r.grad.zero()
# forward, backward, without update
loss_u.forward(clear_no_need_grad=True)
loss_u.backward(clear_buffer=True)
# Normalize gradinet vector and input to variable
r.d = get_direction(r.g)
##### Calculate loss for unlabeled data #####
# Clear remained gradients
solver.zero_grad()
# Set epsilon, the adversarial noise scaling parameter.
eps.data.fill(args.eps_for_vat)
# forward, backward and update
loss_u.forward(clear_no_need_grad=True)
loss_u.backward(clear_buffer=True)
solver.weight_decay(args.weight_decay)
solver.update()
##### Learning rate update #####
if i % args.iter_per_epoch == 0:
solver.set_learning_rate(
solver.learning_rate() * args.learning_rate_decay)
monitor_time.add(i)
# Evaluate the final model by the error rate with validation dataset
valid_error = calc_validation_error(di_v, xv, tv, hv, args.val_iter)
monitor_verr.add(i, valid_error)
monitor_time.add(i)
# Save the model.
nnp_file = os.path.join(
args.model_save_path, 'vat_%06d.nnp' % args.max_iter)
runtime_contents = {
'networks': [
{'name': 'Validation',
'batch_size': args.batchsize_v,
'outputs': {'y': hv},
'names': {'x': xv}}],
'executors': [
{'name': 'Runtime',
'network': 'Validation',
'data': ['x'],
'output': ['y']}]}
save.save(nnp_file, runtime_contents)
from cpp_forward_check import check_cpp_forward
check_cpp_forward(args.model_save_path, [xv.d], [xv], hv, nnp_file)
nnabla_config.set('DATA_ITERATOR', 'data_source_file_cache_size',
'{}'.format(args.cache_size))
nnabla_config.set('DATA_ITERATOR', 'cache_file_format', '.h5')
HERE = os.path.dirname(__file__)
nnabla_examples_root = os.path.join(HERE, '../../../../nnabla-examples')
mnist_examples_root = os.path.realpath(
os.path.join(nnabla_examples_root, 'mnist-collection'))
sys.path.append(mnist_examples_root)
from mnist_data import MnistDataSource
mnist_training_cache = args.output + '/mnist_training.cache'
if not os.path.exists(mnist_training_cache):
os.makedirs(mnist_training_cache)
DataSourceWithFileCache(data_source=MnistDataSource(train=True,
shuffle=False,
rng=None),
cache_dir=mnist_training_cache,
shuffle=False,
rng=None)
mnist_test_cache = args.output + '/mnist_test.cache'
if not os.path.exists(mnist_test_cache):
os.makedirs(mnist_test_cache)
DataSourceWithFileCache(data_source=MnistDataSource(train=False,
shuffle=False,
rng=None),
cache_dir=mnist_test_cache,
shuffle=False,
rng=None)