def monitor_layers_wms(workspace=None):
"""Utility script to download the largest extent for all published WMS layers.
"""
if not workspace:
workspace = os.getenv('GEOSERVER_WORKSPACE')
LOGGER.info('Querying for published layers')
layers = get_layers(workspace)
success = 0
failures = []
# Query each WMS extent:
for layer in layers.keys():
r = layer_getmap_extent(workspace, layer)
if r.headers['Content-Type'] == 'image/jpeg':
LOGGER.info('Queried {}'.format(layer))
success += 1
else:
LOGGER.warning('Failed to query {}'.format(layer))
failures.append(layer)
time.sleep(2) # Pause 2 seconds between requests.
LOGGER.info('{}/{} published layers successfully queried'.format(
success, len(layers)))
if failures:
LOGGER.info('Failed layers: {}'.format(', '.join(failures)))
def _set_model(self):
tf.keras.backend.clear_session()
self.layers = utils.get_layers(self.params['model']['layer'])
self.model = hp_nn.HealpyGCNN(nside=self.params['model']['nside'],
indices=self.indices_ext,
layers=self.layers)
self.model.build(
input_shape=(self.params['dataloader']['batch_size'],
self.pixel_num,
self.params['dataloader']['tomographic_bin_number']))
logger.debug(
f"Building model with input shape ({self.params['dataloader']['batch_size'], self.pixel_num, self.params['dataloader']['tomographic_bin_number']})"
)
if self.params['model']['checkpoint_dir'] == "undefined":
logger.critical(
"Please define the directory within NGSFweights containing the desired weights "
+
"E.g. --checkpoint_dir=layer_2/pixel_noise/01-14-2021-18-38" +
self.worker_id)
sys.exit(0)
else:
path_to_weights = os.path.join(
self.params['model']['weights_dir'],
self.params['model']['checkpoint_dir'])
self.model.load_weights(
tf.train.latest_checkpoint(path_to_weights))
def mp_handler(cddp_path=None):
"""Multiprocessing handler to import metadata from file GDBs in the mounted CDDP volume.
"""
if not cddp_path:
# Assume that this path set via an environment variable if not explicitly passed in.
cddp_path = os.getenv('CDDP_PATH')
datasets = parse_cddp_qmls(cddp_path, LOGGER)
workspace = os.getenv('GEOSERVER_WORKSPACE')
layers = get_layers(workspace)
LOGGER.info('{} datasets scheduled for metadata & style updates'.format(len(datasets)))
# Use a multiprocessing Pool to update layer metadata in parallel.
p = Pool(processes=4)
iterable = [(dataset, layers) for dataset in datasets]
p.starmap(update_metadata, iterable)
z0 = 50
# Sources/Receivers coordinates
s0, ns, ds = x_min + 500, 201, 150 # Sources
r0, nr, dr = x_min + 500, 201, 150 # Receivers
s = [np.arange(s0, s0 + ns * ds, ds), (y_max - z0) * np.ones(ns)] # [sx, sy]
r = [np.arange(r0, r0 + nr * dr, dr), (y_max - z0) * np.ones(nr)] # [rx, ry]
max_frequency = 40.0 # Max freq to resolve - elements/wavelength
# Model
vp = [0.0, 2110.0, 2179.0, 2022.0, 1918.0, 2385.0, 1760.0, 2259.0]
rho = [2000.0 for i in range(len(vp))]
depth = [0.0, 500.0, 700.0, 900.0, 1250.0, 1500.0, 1800.0, 2000.0]
layered_model = ut.get_layers(vp=vp,
rho=rho,
depth=depth,
x_max=4000.0,
nsmooth=5)
# ------------------------------------------------------------------------------
# CREATE NEW SALVUS PROJECT
# ------------------------------------------------------------------------------
# !rm -rf salvus_project
# This line should not be here!
vm = sn.model.volume.cartesian.GenericModel(name="layered_model",
data=layered_model)
p = sn.Project.from_volume_model(path="salvus_project", volume_model=vm)
# Hello
# %%
# ------------------------------------------------------------------------------
def main(args):
"""main
:param args:
argparse.Namespace object from argparse.parse_args().
"""
# Unpack command-line arguments.
train_dir = args.train_dir
style_img_path = args.style_img_path
model_name = args.model_name
preprocess_size = args.preprocess_size
batch_size = args.batch_size
n_epochs = args.n_epochs
run_name = args.run_name
learn_rate = args.learn_rate
loss_content_layers = args.loss_content_layers
loss_style_layers = args.loss_style_layers
content_weights = args.content_weights
style_weights = args.style_weights
num_steps_ckpt = args.num_steps_ckpt
num_pipe_buffer = args.num_pipe_buffer
num_steps_break = args.num_steps_break
beta_val = args.beta
style_target_resize = args.style_target_resize
upsample_method = args.upsample_method
# Load in style image that will define the model.
style_img = utils.imread(style_img_path)
style_img = utils.imresize(style_img, style_target_resize)
style_img = style_img[np.newaxis, :].astype(np.float32)
# Alter the names to include a namescope that we'll use + output suffix.
loss_style_layers = ['vgg/' + i + ':0' for i in loss_style_layers]
loss_content_layers = ['vgg/' + i + ':0' for i in loss_content_layers]
# Get target Gram matrices from the style image.
with tf.variable_scope('vgg'):
X_vgg = tf.placeholder(tf.float32, shape=style_img.shape, name='input')
vggnet = vgg16.vgg16(X_vgg)
with tf.Session() as sess:
vggnet.load_weights('libs/vgg16_weights.npz', sess)
print('Precomputing target style layers.')
target_grams = sess.run(utils.get_grams(loss_style_layers),
feed_dict={X_vgg: style_img})
# Clean up so we can re-create vgg connected to our image network.
print('Resetting default graph.')
tf.reset_default_graph()
# Load in image transformation network into default graph.
shape = [batch_size] + preprocess_size + [3]
with tf.variable_scope('img_t_net'):
X = tf.placeholder(tf.float32, shape=shape, name='input')
Y = create_net(X, upsample_method)
# Connect vgg directly to the image transformation network.
with tf.variable_scope('vgg'):
vggnet = vgg16.vgg16(Y)
# Get the gram matrices' tensors for the style loss features.
input_img_grams = utils.get_grams(loss_style_layers)
# Get the tensors for content loss features.
content_layers = utils.get_layers(loss_content_layers)
# Create loss function
content_targets = tuple(
tf.placeholder(tf.float32,
shape=layer.get_shape(),
name='content_input_{}'.format(i))
for i, layer in enumerate(content_layers))
cont_loss = losses.content_loss(content_layers, content_targets,
content_weights)
style_loss = losses.style_loss(input_img_grams, target_grams,
style_weights)
tv_loss = losses.tv_loss(Y)
beta = tf.placeholder(tf.float32, shape=[], name='tv_scale')
loss = cont_loss + style_loss + beta * tv_loss
with tf.name_scope('summaries'):
tf.summary.scalar('loss', loss)
tf.summary.scalar('style_loss', style_loss)
tf.summary.scalar('content_loss', cont_loss)
tf.summary.scalar('tv_loss', beta * tv_loss)
# Setup input pipeline (delegate it to CPU to let GPU handle neural net)
files = tf.train.match_filenames_once(train_dir + '/train-*')
with tf.variable_scope('input_pipe'), tf.device('/cpu:0'):
batch_op = datapipe.batcher(files, batch_size, preprocess_size,
n_epochs, num_pipe_buffer)
# We do not want to train VGG, so we must grab the subset.
train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='img_t_net')
# Setup step + optimizer
global_step = tf.Variable(0, name='global_step', trainable=False)
optimizer = tf.train.AdamOptimizer(learn_rate).minimize(
loss, global_step, train_vars)
# Setup subdirectory for this run's Tensoboard logs.
if not os.path.exists('./summaries/train/'):
os.makedirs('./summaries/train/')
if run_name is None:
current_dirs = [
name for name in os.listdir('./summaries/train/')
if os.path.isdir('./summaries/train/' + name)
]
name = model_name + '0'
count = 0
while name in current_dirs:
count += 1
name = model_name + '{}'.format(count)
run_name = name
# Savers and summary writers
if not os.path.exists('./training'): # Dir that we'll later save .ckpts to
os.makedirs('./training')
if not os.path.exists('./models'): # Dir that save final models to
os.makedirs('./models')
saver = tf.train.Saver()
final_saver = tf.train.Saver(train_vars)
merged = tf.summary.merge_all()
full_log_path = './summaries/train/' + run_name
train_writer = tf.summary.FileWriter(full_log_path)
# We must include local variables because of batch pipeline.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
# Begin training.
print('Starting training...')
with tf.Session() as sess:
# Initialization
sess.run(init_op)
vggnet.load_weights('libs/vgg16_weights.npz', sess)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
try:
while not coord.should_stop():
current_step = sess.run(global_step)
batch = sess.run(batch_op)
# Collect content targets
content_data = sess.run(content_layers, feed_dict={Y: batch})
feed_dict = {
X: batch,
content_targets: content_data,
beta: beta_val
}
if (current_step % num_steps_ckpt == 0):
# Save a checkpoint
save_path = 'training/' + model_name + '.ckpt'
saver.save(sess, save_path, global_step=global_step)
summary, _, loss_out = sess.run([merged, optimizer, loss],
feed_dict=feed_dict)
train_writer.add_summary(summary, current_step)
print(current_step, loss_out)
elif (current_step % 10 == 0):
# Collect some diagnostic data for Tensorboard.
summary, _, loss_out = sess.run([merged, optimizer, loss],
feed_dict=feed_dict)
train_writer.add_summary(summary, current_step)
# Do some standard output.
print(current_step, loss_out)
else:
_, loss_out = sess.run([optimizer, loss],
feed_dict=feed_dict)
# Throw error if we reach number of steps to break after.
if current_step == num_steps_break:
print('Done training.')
break
except tf.errors.OutOfRangeError:
print('Done training.')
finally:
# Save the model (the image transformation network) for later usage
# in predict.py
final_saver.save(sess, 'models/' + model_name + '_final.ckpt')
coord.request_stop()
coord.join(threads)
def style_transfer(c='./db/pikachu.jpg',
s='./db/starry.jpg',
savename='pikachu-starry',
G_pretrained=None,
epochs=10000,
c_layer=5,
alpha=1,
beta=1e4,
printevery=500,
starting=0):
print(f"Content Image:{c} | Style Image:{s} | savename: {savename}",
flush=True)
# load model
model = models.vgg19(pretrained=True)
model = model.cuda()
for param in model.parameters():
param.requires_grad = False
# load image
# contentImage
img = cv2.imread(c)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
h, w, c = img.shape
contentImage = torch.tensor(img / 255.0).float().cuda()
# style image
img = cv2.imread(s)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (w, h))
styleImage = torch.tensor(img / 255.0).float().cuda()
layers = utils.get_layers(model)
aCs = utils.get_feature_maps(contentImage, layers)
aSs = utils.get_feature_maps(styleImage, layers)
if G_pretrained != 'None':
G = G_pretrained
else:
# torch.manual_seed(0)
# G = torch.rand(contentImage.shape, requires_grad=True, device="cuda")
G = contentImage.detach().clone().requires_grad_(True).cuda()
style_layer_weights = [1.0 / 16 for i in range(16)]
optimizer = optim.AdamW([G], 0.001)
if not os.path.exists(f'./generated/{savename}'):
os.mkdir(f'./generated/{savename}')
# learn stlye + contents
for it in range(starting, starting + epochs):
optimizer.zero_grad()
aGs = utils.get_feature_maps(G, layers)
loss, content_cost, style_cost = utils.compute_total_cost(
aGs,
aCs,
aSs,
style_layer_weights,
content_layer_idx=c_layer,
alpha=alpha,
beta=beta)
if (it + 1) % printevery == 0 or it == 0:
print(
f'iters: {it+1:5d} | loss:{loss.data.cpu().item():2.3e} | content: {content_cost.item():2.3e} | style_cost:{style_cost.item():2.3e}',
flush=True)
save_image(G.permute(2, 0, 1).cpu().detach(),
fp='./generated/{}/iter_{}.png'.format(
savename, it + 1))
loss.backward()
optimizer.step()
def main(args):
# Unpack command-line arguments.
style_img_path = args.style_img_path
cont_img_path = args.cont_img_path
learn_rate = args.learn_rate
loss_content_layers = args.loss_content_layers
loss_style_layers = args.loss_style_layers
content_weights = args.content_weights
style_weights = args.style_weights
num_steps_break = args.num_steps_break
beta = args.beta
style_target_resize = args.style_target_resize
cont_target_resize = args.cont_target_resize
output_img_path = args.output_img_path
# Load in style image that will define the model.
style_img = utils.imread(style_img_path)
style_img = utils.imresize(style_img, style_target_resize)
style_img = style_img[np.newaxis, :].astype(np.float32)
# Alter the names to include a namescope that we'll use + output suffix.
loss_style_layers = ['vgg/' + i + ':0' for i in loss_style_layers]
loss_content_layers = ['vgg/' + i + ':0' for i in loss_content_layers]
# Get target Gram matrices from the style image.
with tf.variable_scope('vgg'):
X_vgg = tf.placeholder(tf.float32, shape=style_img.shape, name='input')
vggnet = vgg16.vgg16(X_vgg)
with tf.Session() as sess:
vggnet.load_weights('libs/vgg16_weights.npz', sess)
print 'Precomputing target style layers.'
target_grams = sess.run(utils.get_grams(loss_style_layers),
feed_dict={'vgg/input:0': style_img})
# Clean up so we can re-create vgg at size of input content image for
# training.
print 'Resetting default graph.'
tf.reset_default_graph()
# Read in + resize the content image.
cont_img = utils.imread(cont_img_path)
cont_img = utils.imresize(cont_img, cont_target_resize)
cont_img = cont_img[np.newaxis, :].astype(np.float32)
# Setup VGG and initialize it with white noise image that we'll optimize.
shape = cont_img.shape
with tf.variable_scope('to_train'):
white_noise = np.random.rand(shape[0], shape[1], shape[2],
shape[3]) * 255.0
white_noise = tf.constant(white_noise.astype(np.float32))
X = tf.get_variable('input', dtype=tf.float32, initializer=white_noise)
with tf.variable_scope('vgg'):
vggnet = vgg16.vgg16(X)
# Get the gram matrices' tensors for the style loss features.
input_img_grams = utils.get_grams(loss_style_layers)
# Get the tensors for content loss features.
content_layers = utils.get_layers(loss_content_layers)
# Get the target content features
with tf.Session() as sess:
vggnet.load_weights('libs/vgg16_weights.npz', sess)
print 'Precomputing target content layers.'
content_targets = sess.run(content_layers,
feed_dict={'to_train/input:0': cont_img})
# Create loss function
cont_loss = losses.content_loss(content_layers, content_targets,
content_weights)
style_loss = losses.style_loss(input_img_grams, target_grams,
style_weights)
tv_loss = losses.tv_loss(X)
loss = cont_loss + style_loss + beta * tv_loss
# We do not want to train VGG, so we must grab the subset.
train_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES,
scope='to_train')
# Setup step + optimizer
global_step = tf.Variable(0, name='global_step', trainable=False)
optimizer = tf.train.AdamOptimizer(learn_rate) \
.minimize(loss, global_step, train_vars)
# Initializer
init_op = tf.global_variables_initializer()
# Begin training
with tf.Session() as sess:
sess.run(init_op)
vggnet.load_weights('libs/vgg16_weights.npz', sess)
current_step = 0
while current_step < num_steps_break:
current_step = sess.run(global_step)
if (current_step % 10 == 0):
# Collect some diagnostic data for Tensorboard.
_, loss_out = sess.run([optimizer, loss])
# Do some standard output.
print current_step, loss_out
else:
# optimizer.minimize(sess)
_, loss_out = sess.run([optimizer, loss])
# Upon finishing, get the X tensor (our image).
img_out = sess.run(X)
# Save it.
img_out = np.squeeze(img_out)
utils.imwrite(output_img_path, img_out)
def train_deepTrunk(dTNet, args, device, stats, train_loader, test_loader):
epoch = 0
lr = args.lr
n_epochs = args.n_epochs
last_gate_layer = -2
## Get a duplicate of the data loaders that will be changed
train_set_spec = MyVisionDataset.from_idx(train_loader.dataset, np.ones_like(train_loader.dataset.targets).astype(bool), train=True)
train_set_spec.set_weights(None)
train_loader_spec = torch.utils.data.DataLoader(train_set_spec, batch_size=args.train_batch,
shuffle=~args.debug, num_workers=train_loader.num_workers, pin_memory=True, drop_last=True)
test_set_spec = MyVisionDataset.from_idx(test_loader.dataset,
np.ones_like(test_loader.dataset.targets).astype(bool), train=False)
test_set_spec.set_weights(None)
test_loader_spec = torch.utils.data.DataLoader(test_set_spec, batch_size=args.test_batch,
shuffle=False, num_workers=train_loader.num_workers, pin_memory=True, drop_last=False)
### train core-network
if args.gate_on_trunk and args.train_trunk:
layers = get_layers(args.train_mode, dTNet.trunk_cnet, n_attack_layers=args.n_attack_layers,
min_layer=-1, base_layers=False,
protected_layers=args.protected_layers)
assert len(layers) > 0
print("training trunk")
trunk_trainability_state = dTNet.trunk_net.get_freeze_state()
epoch = train_episode(dTNet.trunk_cnet, dTNet, device, args, lr, epoch, n_epochs, train_loader_spec,
test_loader_spec, -1, layers, stats=stats, eps_init=0)
### Start training of all certification networks
for k, exit_idx in enumerate(dTNet.exit_ids[1::]):
if (not args.gate_on_trunk) and train_loader_spec is not None and len(train_loader_spec) > 0 and args.load_branch_model is None:
layers = get_layers(args.train_mode, dTNet.branch_cnets[exit_idx],
n_attack_layers=args.n_attack_layers,
min_layer=-1, base_layers=False,
protected_layers=args.protected_layers)
print("Training branch %d" % exit_idx)
net_weights = [1] if args.cotrain_entropy is None else [1 - args.cotrain_entropy, args.cotrain_entropy]
branch_trainability_state = dTNet.branch_nets[exit_idx].get_freeze_state()
epoch = train_episode(dTNet.branch_cnets[exit_idx]
if not (dTNet.gate_type == "entropy" and args.cotrain_entropy is not None)
else
[dTNet.branch_cnets[exit_idx], dTNet.gate_cnets[exit_idx]], dTNet, device, args,
lr, epoch, n_epochs, train_loader_spec, test_loader_spec, exit_idx, layers,
stats=stats, net_weights=net_weights)
if args.retrain_branch:
dTNet.branch_nets[exit_idx].restore_freeze(branch_trainability_state)
### train gate unless entropy is used for selection
if args.gate_type == "net":
print("getting dataset for gate training")
gate_train_loader, gate_test_loader = get_gated_data_loaders(args, device,
dTNet.trunk_cnet if args.gate_on_trunk else dTNet.branch_cnets[exit_idx],
train_loader_spec, test_loader_spec, stats,
mode=args.gate_mode,
exact=args.exact_target_cert)
min_layer = -2
if args.gate_feature_extraction is not None:
extraction_layer = [ii for ii in range(len(dTNet.gate_nets[exit_idx].blocks)) if isinstance(dTNet.gate_nets[exit_idx].blocks[ii], Linear)]
extraction_layer = extraction_layer[-min(len(extraction_layer), args.gate_feature_extraction)]
min_layer = [ii for ii in range(len(dTNet.gate_nets[exit_idx].blocks)) if isinstance(dTNet.gate_nets[exit_idx].blocks[ii], ReLU) and ii<extraction_layer][-1]
layers = get_layers(args.train_mode, dTNet.gate_cnets[exit_idx], n_attack_layers=args.n_attack_layers,
protected_layers=args.protected_layers, min_layer=min_layer)
print("Training gate %d" % (exit_idx))
epoch = train_episode(dTNet.gate_cnets[exit_idx], dTNet, device, args, lr, epoch, n_epochs,
gate_train_loader, gate_test_loader, exit_idx, layers, stats,
balanced_loss=args.balanced_gate_loss)
### train certification-network
if args.retrain_branch:
print("Getting datasets branch %d" % exit_idx)
branch_train_loader, branch_test_loader = get_gated_data_loaders(args, device,
dTNet.gate_cnets[exit_idx],
train_loader_spec,
test_loader_spec,
stats, mode="branch",
exact=args.exact_target_cert)
layers = get_layers(args.train_mode, dTNet.branch_cnets[exit_idx],
n_attack_layers=args.n_attack_layers,
min_layer=-1, base_layers=False,
protected_layers=args.protected_layers)
print("Training branch for gate %d" % exit_idx)
net_weights = [1] if args.cotrain_entropy is None else [1 - args.cotrain_entropy, args.cotrain_entropy]
if len(branch_train_loader) > 0:
epoch = train_episode(dTNet.branch_cnets[exit_idx]
if not (dTNet.gate_type == "entropy" and args.cotrain_entropy is not None)
else [dTNet.branch_cnets[exit_idx], dTNet.gate_cnets[exit_idx]],
dTNet, device, args, lr, epoch, n_epochs, branch_train_loader, branch_test_loader,
exit_idx, layers, stats=stats,
eps_init=0 if args.gate_on_trunk else None, net_weights=net_weights)
else:
raise RuntimeWarning("Certification-Network not reached by any training samples. Training skipped.")
if args.train_trunk and (args.retrain_trunk or not args.gate_on_trunk):
train_loader_spec, test_loader_spec = get_gated_data_loaders(args, device, dTNet.gate_cnets[exit_idx],
train_loader_spec,
test_loader_spec,
stats, mode="trunk",
threshold=0 if args.gate_type == "net" else dTNet.threshold[exit_idx],
exact=args.exact_target_cert)
### Train the core network (again)
if args.train_trunk and (args.retrain_trunk or not args.gate_on_trunk):
if train_loader_spec is not None and len(train_loader_spec) > 0:
print("training trunk final")
dTNet.trunk_net.restore_freeze(trunk_trainability_state)
layers = get_layers(args.train_mode, dTNet.trunk_cnet, n_attack_layers=args.n_attack_layers,
min_layer=last_gate_layer, protected_layers=args.protected_layers, base_layers=False)
epoch = train_episode(dTNet.trunk_cnet, dTNet, device, args, lr, epoch, n_epochs, train_loader_spec,
test_loader_spec, -1, layers, stats=stats, eps_init=None if (
(args.retrain_trunk and args.gate_on_trunk) or (
args.load_model is None and args.gate_type == "net")) else 0)
else:
raise RuntimeWarning("Core-Network not reached by any training samples. Training skipped.")
return epoch
def main(unused_agrv=None):
"""main
:param args:
argparse.Namespace object from argparse.parse_args().
"""
# Unpack command-line arguments.
train_dir = FLAGS.train_dir
style_dataset = FLAGS.style_dataset
model_name = FLAGS.model_name
preprocess_size = [FLAGS.image_size, FLAGS.image_size]
batch_size = FLAGS.batch_size
n_epochs = FLAGS.n_epochs
run_name = FLAGS.run_name
checkpoint = FLAGS.checkpoint
learn_rate = FLAGS.learning_rate
content_weights = FLAGS.content_weights
style_weights = FLAGS.style_weights
num_pipe_buffer = FLAGS.num_pipe_buffer
style_coefficients = FLAGS.style_coefficients
num_styles = FLAGS.num_styles
train_steps = FLAGS.train_steps
upsample_method = FLAGS.upsample_method
# Setup input pipeline (delegate it to CPU to let GPU handle neural net)
files = tf.train.match_filenames_once(train_dir + '/train-*')
style_files = tf.train.match_filenames_once(style_dataset)
print("style %s" % style_files)
with tf.variable_scope('input_pipe'), tf.device('/cpu:0'):
_, style_labels, style_grams = datapipe.style_batcher(
style_files, batch_size, preprocess_size, n_epochs,
num_pipe_buffer)
batch_op = datapipe.batcher(files, batch_size, preprocess_size,
n_epochs, num_pipe_buffer)
""" Set up the style coefficients """
if style_coefficients is None:
style_coefficients = [1.0 for _ in range(num_styles)]
else:
style_coefficients = ast.literal_eval(style_coefficients)
if len(style_coefficients) != num_styles:
raise ValueError(
'number of style coeffients differs from number of styles')
style_coefficient = tf.gather(tf.constant(style_coefficients),
style_labels)
""" Set up weight of style and content image """
content_weights = ast.literal_eval(content_weights)
style_weights = ast.literal_eval(style_weights)
style_weights = dict([(key, style_coefficient * val)
for key, val in style_weights.iteritems()])
target_grams = []
for name, val in style_weights.iteritems():
target_grams.append(style_grams[name])
# Alter the names to include a name_scope that we'll use + output suffix.
loss_style_layers = []
loss_style_weights = []
loss_content_layers = []
loss_content_weights = []
for key, val in style_weights.iteritems():
loss_style_layers.append(key + ':0')
loss_style_weights.append(val)
for key, val in content_weights.iteritems():
loss_content_layers.append(key + ':0')
loss_content_weights.append(val)
# Load in image transformation network into default graph.
shape = [batch_size] + preprocess_size + [3]
with tf.variable_scope('styleNet'):
X = tf.placeholder(tf.float32, shape=shape, name='input')
Y = transform(X, style_labels, num_styles, upsample_method)
print(Y)
# Connect vgg directly to the image transformation network.
with tf.variable_scope('vgg'):
vggnet = vgg16.vgg16(Y)
# Get the gram matrices' tensors for the style loss features.
input_img_grams = utils.get_grams(loss_style_layers)
# Get the tensors for content loss features.
content_layers = utils.get_layers(loss_content_layers)
# Create loss function
content_targets = tuple(
tf.placeholder(tf.float32,
shape=layer.get_shape(),
name='content_input_{}'.format(i))
for i, layer in enumerate(content_layers))
cont_loss = losses.content_loss(content_layers, content_targets,
loss_content_weights)
style_loss = losses.style_loss(input_img_grams, target_grams,
loss_style_weights)
loss = cont_loss + style_loss
with tf.name_scope('summaries'):
tf.summary.scalar('loss', loss)
tf.summary.scalar('style_loss', style_loss)
tf.summary.scalar('content_loss', cont_loss)
# We do not want to train VGG, so we must grab the subset.
other_vars = [
var for var in tf.get_variable_scope('styleNet')
if 'CondInstNorm' not in var.name
]
train_vars = [
var for var in tf.get_variable_scope('styleNet')
if 'CondInstNorm' in var.name
]
# Setup step + optimizer
global_step = tf.Variable(0, name='global_step', trainable=False)
optimizer = tf.train.AdamOptimizer(learn_rate).minimize(
loss, global_step, train_vars)
# Setup subdirectory for this run's Tensoboard logs.
if not os.path.exists('./summaries/train/'):
os.makedirs('./summaries/train/')
if run_name is None:
current_dirs = [
name for name in os.listdir('./summaries/train/')
if os.path.isdir('./summaries/train/' + name)
]
name = model_name + '0'
count = 0
while name in current_dirs:
count += 1
name = model_name + '{}'.format(count)
run_name = name
# Savers and summary writers
if not os.path.exists('./training'): # Dir that we'll later save .ckpts to
os.makedirs('./training')
if not os.path.exists('./models'): # Dir that save final models to
os.makedirs('./models')
saver = tf.train.Saver()
saver_n_stylee = tf.train.Saver(other_vars)
final_saver = tf.train.Saver(train_vars)
merged = tf.summary.merge_all()
full_log_path = './summaries/train/' + run_name
train_writer = tf.summary.FileWriter(full_log_path, tf.Session().graph)
# We must include local variables because of batch pipeline.
init_op = tf.group(tf.global_variables_initializer(),
tf.local_variables_initializer())
# Begin training.
print 'Starting training...'
with tf.Session() as sess:
# Initialization
sess.run(init_op)
vggnet.load_weights(vgg16.checkpoint_file(), sess)
saver_n_stylee.restore(sess, checkpoint)
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
try:
while not coord.should_stop():
current_step = sess.run(global_step)
batch = sess.run(batch_op)
# Collect content targets
content_data = sess.run(content_layers, feed_dict={Y: batch})
feed_dict = {X: batch, content_targets: content_data}
if (current_step % 1000 == 0):
# Save a checkpoint
save_path = 'training/' + model_name + '.ckpt'
saver.save(sess, save_path, global_step=global_step)
summary, _, loss_out, c_loss, s_loss = sess.run(
[merged, optimizer, loss, cont_loss, style_loss],
feed_dict=feed_dict)
train_writer.add_summary(summary, current_step)
print current_step, loss_out, c_loss, s_loss
elif (current_step % 10 == 0):
# Collect some diagnostic data for Tensorboard.
summary, _, loss_out, c_loss, s_loss = sess.run(
[merged, optimizer, loss, cont_loss, style_loss],
feed_dict=feed_dict)
train_writer.add_summary(summary, current_step)
# Do some standard output.
# if (current_step % 1000 == 0):
print current_step, loss_out, c_loss, s_loss
else:
_, loss_out = sess.run([optimizer, loss],
feed_dict=feed_dict)
# Throw error if we reach number of steps to break after.
if current_step == train_steps:
print('Done training.')
break
except tf.errors.OutOfRangeError:
print('Done training.')
finally:
# Save the model (the image transformation network) for later usage
# in predict.py
final_saver.save(sess,
'models/' + model_name + '_final.ckpt',
write_meta_graph=False)
coord.request_stop()
coord.join(threads)
def run(args=None):
device = 'cuda' if torch.cuda.is_available() and (
not args.no_cuda) else 'cpu'
num_train, train_loader, test_loader, input_size, input_channel, n_class = get_loaders(
args)
lossFn = nn.CrossEntropyLoss(reduction='none')
evalFn = lambda x: torch.max(x, dim=1)[1]
net = get_net(device,
args.dataset,
args.net,
input_size,
input_channel,
n_class,
load_model=args.load_model,
net_dim=args.cert_net_dim
) #, feature_extract=args.core_feature_extract)
timestamp = int(time.time())
model_signature = '%s/%s/%d/%s_%.5f/%d' % (args.dataset, args.exp_name,
args.exp_id, args.net,
args.train_eps, timestamp)
model_dir = args.root_dir + 'models_new/%s' % (model_signature)
args.model_dir = model_dir
count_vars(args, net)
if not os.path.exists(model_dir):
os.makedirs(model_dir)
if isinstance(net, UpscaleNet):
relaxed_net = None
relu_ids = None
else:
relaxed_net = RelaxedNetwork(net.blocks, args.n_rand_proj).to(device)
relu_ids = relaxed_net.get_relu_ids()
if "nat" in args.train_mode:
cnet = CombinedNetwork(net,
relaxed_net,
lossFn=lossFn,
evalFn=evalFn,
device=device,
no_r_net=True).to(device)
else:
dummy_input = torch.rand((1, ) + net.dims[0],
device=device,
dtype=torch.float32)
cnet = CombinedNetwork(net,
relaxed_net,
lossFn=lossFn,
evalFn=evalFn,
device=device,
dummy_input=dummy_input).to(device)
n_epochs, test_nat_loss, test_nat_acc, test_adv_loss, test_adv_acc = args.n_epochs, None, None, None, None
if 'train' in args.train_mode:
tb_writer = SummaryWriter(model_dir)
stats = Statistics(len(train_loader), tb_writer, model_dir)
args_file = os.path.join(model_dir, 'args.json')
with open(args_file, 'w') as fou:
json.dump(vars(args), fou, indent=4)
write_config(args, os.path.join(model_dir, 'run_config.txt'))
eps = 0
epoch = 0
lr = args.lr
n_epochs = args.n_epochs
if "COLT" in args.train_mode:
relu_stable = args.relu_stable
# if args.layers is None:
# args.layers = [-2, -1] + relu_ids
layers = get_layers(args.train_mode,
cnet,
n_attack_layers=args.n_attack_layers,
protected_layers=args.protected_layers)
elif "adv" in args.train_mode:
relu_stable = None
layers = [-1, -1]
args.mix = False
elif "natural" in args.train_mode:
relu_stable = None
layers = [-2, -2]
args.nat_factor = 1
args.mix = False
elif "diffAI" in args.train_mode:
relu_stable = None
layers = [-2, -2]
else:
assert False, "Unknown train mode %s" % args.train_mode
print('Saving model to:', model_dir)
print('Training layers: ', layers)
for j in range(len(layers) - 1):
opt, lr_scheduler = get_opt(cnet.net,
args.opt,
lr,
args.lr_step,
args.lr_factor,
args.n_epochs,
train_loader,
args.lr_sched,
fixup="fixup" in args.net)
curr_layer_idx = layers[j + 1]
eps_old = eps
eps = get_scaled_eps(args, layers, relu_ids, curr_layer_idx, j)
kappa_sched = Scheduler(0.0 if args.mix else 1.0, 1.0,
num_train * args.mix_epochs, 0)
beta_sched = Scheduler(
args.beta_start if args.mix else args.beta_end, args.beta_end,
args.train_batch * len(train_loader) * args.mix_epochs, 0)
eps_sched = Scheduler(eps_old if args.anneal else eps, eps,
num_train * args.anneal_epochs, 0)
layer_dir = '{}/{}'.format(model_dir, curr_layer_idx)
if not os.path.exists(layer_dir):
os.makedirs(layer_dir)
print('\nnew train phase: eps={:.5f}, lr={:.2e}, curr_layer={}\n'.
format(eps, lr, curr_layer_idx))
for curr_epoch in range(n_epochs):
train(device,
epoch,
args,
j + 1,
layers,
cnet,
eps_sched,
kappa_sched,
opt,
train_loader,
lr_scheduler,
relu_ids,
stats,
relu_stable,
relu_stable_protected=args.relu_stable_protected,
beta_sched=beta_sched)
if isinstance(lr_scheduler, optim.lr_scheduler.StepLR
) and curr_epoch >= args.mix_epochs:
lr_scheduler.step()
if (epoch + 1) % args.test_freq == 0:
with torch.no_grad():
test_nat_loss, test_nat_acc, test_adv_loss, test_adv_acc = test(
device,
args,
cnet,
test_loader if args.test_set == "test" else
train_loader, [curr_layer_idx],
stats=stats,
log_ind=(epoch + 1) % n_epochs == 0)
if (epoch + 1) % args.test_freq == 0 or (epoch +
1) % n_epochs == 0:
torch.save(
net.state_dict(),
os.path.join(layer_dir, 'net_%d.pt' % (epoch + 1)))
torch.save(
opt.state_dict(),
os.path.join(layer_dir, 'opt_%d.pt' % (epoch + 1)))
stats.update_tb(epoch)
epoch += 1
relu_stable = None if relu_stable is None else relu_stable * args.relu_stable_layer_dec
lr = lr * args.lr_layer_dec
if args.cert:
with torch.no_grad():
diffAI_cert(
device,
args,
cnet,
test_loader if args.test_set == "test" else train_loader,
stats=stats,
log_ind=True,
epoch=epoch,
domains=args.cert_domain)
elif args.train_mode == 'print':
print('printing network to:', args.out_net_file)
dummy_input = torch.randn(1,
input_channel,
input_size,
input_size,
device='cuda')
net.skip_norm = True
torch.onnx.export(net, dummy_input, args.out_net_file, verbose=True)
elif args.train_mode == 'test':
with torch.no_grad():
test(device,
args,
cnet,
test_loader if args.test_set == "test" else train_loader,
[-1],
log_ind=True)
elif args.train_mode == "cert":
tb_writer = SummaryWriter(model_dir)
stats = Statistics(len(train_loader), tb_writer, model_dir)
args_file = os.path.join(model_dir, 'args.json')
with open(args_file, 'w') as fou:
json.dump(vars(args), fou, indent=4)
write_config(args, os.path.join(model_dir, 'run_config.txt'))
print('Saving results to:', model_dir)
with torch.no_grad():
diffAI_cert(
device,
args,
cnet,
test_loader if args.test_set == "test" else train_loader,
stats=stats,
log_ind=True,
domains=args.cert_domain)
exit(0)
else:
assert False, 'Unknown mode: {}!'.format(args.train_mode)
return test_nat_loss, test_nat_acc, test_adv_loss, test_adv_acc
