def put_file_in_bucket(bucket_name):
# get file name to upload to bucket
file_path = utils.get_abs_file_path('\nPath of file to put into bucket: ')
try: # set object name to just the base name as would otherwise create the folders to file.
# Would have problems, especially when path to object is relative and when getting file link
response = s3.Object(bucket_name, os.path.basename(file_path)).put(
ACL='public-read', # make public readable
Body=open(file_path, 'rb'))
utils.print_and_log(response)
choice = input("Would you like to append file to index html (y/n): ").lower()
if choice == 'y':
url = get_file_url(bucket_name, os.path.basename(file_path))
public_ip = get_instance_ip()
if public_ip: # If there is a public ip returned
while True:
try:
key_path = make_key_read_only(
utils.get_valid_key("Enter path to your private key: "))
append_to_index(public_ip, url, key_path)
break
except Exception as error:
print("Error: Not a valid option")
except Exception as error:
utils.print_and_log(error)
def run_exploit(exploit_dir, container_name, timeout, log=None):
log = print_and_log("[*] Running exploit", log)
script = os.path.join(base_dir(), "launch_exploit.sh")
cmdline = \
"%s %s %s %d %d" % \
(script, container_name, SERVICE_IP, SERVICE_PORT, timeout)
output, err, e = run_command(cmdline, exploit_dir)
if log is not None:
try:
log = log.decode('utf_8', 'ignore')
output = output.decode('utf-8', 'ignore')
log = log + output
except UnicodeDecodeError:
pass
if e != 0:
log = print_and_log("[*] Failed to run exploit", log)
log = print_and_log(err, log)
log = print_and_log("==========================", log)
return None, log
# Exploit prints out the flag string at the end.
tokens = filter(None, output.split('\n')) # Filter out empty strings
flag_candidate = filter(None, tokens)[-1] # Read the last line
return flag_candidate, log
def internal_validate(self,
sess,
train_data,
test_idx_array,
tidx,
test_batch_size=10000):
print_and_log("...Start testing...", self.log_file)
batches_test = batch_gen(train_data[0], train_data[1],
test_idx_array, test_batch_size,
np.array([tidx]), False, True)
tars = np.array([], dtype=np.int32)
scores = np.array([], dtype=np.float32)
while True:
batch_fp, batch_tar, is_last = next(batches_test)
batch_cls = sess.run(self.cls_tensor,
feed_dict={
self.fp_tensor: batch_fp,
self.tar_tensor: batch_tar,
self.idropout: 0,
self.idropout: 0
})
clear_point = np.where(batch_tar != 0)
tars = np.concatenate((tars, batch_tar[clear_point]))
scores = np.concatenate((scores, batch_cls[clear_point]))
if is_last:
break
return tars, scores
def get_instance_usage(key_path, public_ip):
usage_cmd = construct_ssh(key_path, public_ip, " 'top -n 1 -b'")
(status, output) = subprocess.getstatusoutput(usage_cmd)
if status == 0:
utils.print_and_log('\n' + output)
time.sleep(5)
else:
utils.print_and_log('\nGet usage failed ' + output)
def verify_exploit(exploit_dir,
service_dir,
branch,
timeout,
config,
encrypt=False,
log=None):
if not os.path.isdir(exploit_dir):
print("[*] Exploit directory '%s' does not exist" % exploit_dir)
return False, log
if not os.path.isdir(service_dir):
print("[*] Service directory '%s' does not exist" % service_dir)
return False, log
# Create random flag value
flag = random_string(10)
# Start the service
service_dirname = get_dirname(service_dir)
service_container_name = "%s-%s" % (service_dirname,
branch.replace('/', '_'))
result, log = start_service(service_dir, branch, service_container_name, \
flag, log=log)
if not result:
return False, log
time.sleep(2)
# Run the exploit
exploit_dirname = get_dirname(exploit_dir)
exploit_container_name = "exploit-%s" % branch.replace('/', '_')
exploit_result, log = run_exploit(exploit_dir, exploit_container_name, \
timeout, log=log)
# Clean up containers
docker_cleanup(service_container_name)
docker_cleanup(exploit_container_name)
log = print_and_log("[*] Exploit returned : %s" % exploit_result, log)
log = print_and_log("[*] Solution flag : %s" % flag, log)
if exploit_result == flag:
print("[*] Exploit worked successfully")
if encrypt:
print("[*] Encrypting the verified exploit")
# Set your own team as target team, and signer is not needed.
target_team = config["player_team"]
encrypted_file = encrypt_exploit(exploit_dir, target_team, config)
if encrypted_file is None:
print("[*] Failed to encrypt exploit")
else:
print("[*] Your exploit is encrypted in %s" % encrypted_file)
print("[*] Now you may commit and push this encrypted exploit "\
"to the corresponding branch of your service repository")
return True, log
else:
log = print_and_log("[*] Exploit returned a wrong flag string", log)
return False, log
def test_print_and_log(self):
print_and_log("Hello from test") # will append and create if log.txt is not present
self.assertTrue(True, os.path.exists(os.path.abspath('log.txt'))) # assert the path exists
log_file = open("log.txt", "r") # open log.txt
temp_list = []
for line in log_file: # assign temp list to be the last line in log.txt
temp_list = (line.split(' - ')) # split line at ' - ' to separate date and message
self.assertEqual("Hello from test", temp_list[1]) # assert message are equal
log_file.close()
def load_from_file(date_str,
coll_names=('votes', 'views', 'comments', 'posts', 'users',
'tags', 'tagrels', 'sequences')):
"""Loads database collections from csvs to dataframes, ensures datetimes load correctly."""
def read_csv(coll_name):
if coll_name in read_dtypes_arg:
dtypes = read_dtypes_arg[coll_name]
else:
dtypes = None
print_and_log('Reading {}'.format(coll_name))
df = pd.read_csv(complete_path_to_file(coll_name), dtype=dtypes)
# read in all datetime types correctly
for dt_col in [
'postedAt', 'createdAt', 'votedAt', 'startTime', 'endTime',
'earliest_comment', 'most_recent_comment', 'earliest_vote',
'most_recent_vote', 'most_recent_post', 'earliest_post',
'most_recent_activity', 'earliest_activity', 'true_earliest',
'curatedAt', 'earliest_view', 'birth'
]:
if dt_col in df.columns:
df.loc[:, dt_col] = pd.to_datetime(df[dt_col])
return df
def complete_path_to_file(coll_name):
return BASE_PATH + '{folder}/{date}/{coll_name}.csv'.format(
folder='processed', date=date_str, coll_name=coll_name)
if date_str == 'most_recent':
date_str = get_list_of_dates()[0][-8:]
read_dtypes_arg = {
'users': None,
'posts': None,
'comments': None,
'votes': {
'collectionName': 'category',
'voteType': 'category',
'afPower': 'int8',
'power': 'int8'
},
'views': None,
'sequences': None,
'tags': None,
'tagrels': None
}
print_and_log("Files to be loaded:")
[print(complete_path_to_file(coll_name)) for coll_name in coll_names]
return {coll_name: read_csv(coll_name) for coll_name in coll_names}
def get_file_url(bucket_name, object_name):
client = boto3.client('s3')
url = client.generate_presigned_url('get_object',
Params={
'Bucket': bucket_name,
'Key': object_name,
},
ExpiresIn=3600)
# split url into a list at '?' as the above generates a url link that will be expired
split = url.split('?')
utils.print_and_log('Url: ' + split[0])
return split[0] # Return the url with no expiry date
def write_collection(coll_name, coll_df, date_str): # (string, df, arg_bundle) -> None
# hardcoded to write to db directory. wonderful hardcoding
# this function really needs some cleanin'
print_and_log('Writing {} to disk.'.format(coll_name))
directory = BASE_PATH + '{folder}/{date}'.format(folder='processed', date=date_str) # vestigial folder structure
pathlib.Path(directory).mkdir(exist_ok=True)
coll_df.to_csv(directory + '/{}.csv'.format(coll_name), index=False)
print_and_log('Writing {} to disk completed.\n'.format(coll_name))
return None
def run(self):
config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True
with tf.compat.v1.Session(config=config) as session:
train_accuracies = []
losses = []
total_iterations = self.args.training_iterations
iteration = self.start_iteration
for task_dict in self.video_loader:
if iteration >= total_iterations:
break
iteration += 1
torch.set_grad_enabled(True)
task_loss, task_accuracy = self.train_task(task_dict)
train_accuracies.append(task_accuracy)
losses.append(task_loss)
# optimize
if ((iteration + 1) % self.args.tasks_per_batch
== 0) or (iteration == (total_iterations - 1)):
self.optimizer.step()
self.optimizer.zero_grad()
self.scheduler.step()
if (iteration + 1) % self.args.print_freq == 0:
# print training stats
print_and_log(
self.logfile,
'Task [{}/{}], Train Loss: {:.7f}, Train Accuracy: {:.7f}'
.format(iteration + 1, total_iterations,
torch.Tensor(losses).mean().item(),
torch.Tensor(train_accuracies).mean().item()))
train_accuracies = []
losses = []
if ((iteration + 1) % self.args.save_freq
== 0) and (iteration + 1) != total_iterations:
self.save_checkpoint(iteration + 1)
if ((iteration + 1) in self.args.test_iters
) and (iteration + 1) != total_iterations:
accuracy_dict = self.test(session)
print(accuracy_dict)
self.test_accuracies.print(self.logfile, accuracy_dict)
# save the final model
torch.save(self.model.state_dict(), self.checkpoint_path_final)
self.logfile.close()
def run_ga_pipeline():
##Get Data
ga_metrics = get_all_metrics()
## Postgres Upload
ga_metrics_pg = copy.deepcopy(ga_metrics)
for df in ga_metrics_pg.values():
df.columns = [col.replace(':', '_') for col in df.columns]
try:
engine = get_pg_engine()
with engine.begin() as conn:
print_and_log('truncating postgres tables')
truncate_or_drop_tables(
['ga_' + name for name in ga_metrics.keys()],
conn=conn,
drop=False)
print_and_log('loading tables into postgres db')
[
bulk_upload_to_pg(df, table_name='ga_' + name, conn=conn)
for name, df in ga_metrics_pg.items()
]
print_and_log('transaction successful!')
except:
print_and_log('transfer failed')
finally:
engine.dispose()
def run_pg_pandas_transfer(
dfs,
tables=('users', 'posts', 'comments', 'votes', 'tags', 'tagrels'),
drop_tables=False,
):
dfs_prepared = prep_frames_for_db(dfs)
# try:
engine = get_pg_engine()
with engine.begin() as conn:
if drop_tables:
print_and_log('dropping postgres tables')
else:
print_and_log('truncating postgres tables')
truncate_or_drop_tables(tables, conn=conn, drop=drop_tables)
if drop_tables:
create_tables(tables, conn)
print_and_log('loading tables into postgres db')
[
bulk_upload_to_pg(dfs_prepared[coll], table_name=coll, conn=conn)
for coll in tables
]
print_and_log('transaction successful!')
# except:
# print_and_log('transfer failed')
# finally:
engine.dispose()
def init_model(model_dir, nb_of_classes):
device = torch.device("cuda:0")
resnext50_32x4d = models.resnext50_32x4d(pretrained=True)
resnext50_32x4d.fc = nn.Linear(resnext50_32x4d.fc.in_features, nb_of_classes)
resnext50_32x4d = resnext50_32x4d.to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(resnext50_32x4d.parameters(), lr=0.001)
utils.print_and_log(model_dir, resnext50_32x4d)
return resnext50_32x4d, criterion, optimizer
def copy_check_webserver(public_ip, key_path):
scp_cmd = 'scp -i ' + key_path + ' check_webserver.py ec2-user@' + public_ip + ':.'
utils.print_and_log('Now trying to copy check_webserver to instance')
(status, output) = subprocess.getstatusoutput(scp_cmd)
if status == 0: # if successfully copied, change permissions to make it executable
utils.print_and_log('Successfully copied check_webserver.py to new instance')
(status, output) = subprocess.getstatusoutput(
construct_ssh(key_path, public_ip, " 'chmod 700 ./check_webserver.py'"))
if status == 0: # if successful, run check webserver
utils.print_and_log('Made check_webserver executable on instance')
run_check_webserver(public_ip, key_path)
else:
utils.print_and_log(str('Failed to change persmissions: ' + output))
else:
utils.print_and_log('Copy check_webserver.py failed :(')
def read_csv(coll_name):
print_and_log('Reading {}'.format(coll_name))
df = pd.read_csv(complete_path_to_file(coll_name), dtype=read_dtypes_arg[coll_name])
# read in all datetime types correctly
for dt_col in ['postedAt', 'createdAt', 'votedAt', 'startTime', 'endTime',
'earliest_comment', 'most_recent_comment', 'earliest_vote', 'most_recent_vote',
'most_recent_post', 'earliest_post', 'most_recent_activity', 'earliest_activity',
'true_earliest', 'curatedAt', 'earliest_view'
]:
if dt_col in df.columns:
df.loc[:, dt_col] = pd.to_datetime(df[dt_col])
return df
def install_python35(key_path, public_ip):
install_cmd = construct_ssh(key_path, public_ip, " 'sudo yum install -y python35'")
utils.print_and_log('Installing python35')
(status, output) = subprocess.getstatusoutput(install_cmd)
if status == 0:
utils.print_and_log('Successfully installed python35')
else:
utils.print_and_log('Failed to install python35')
utils.print_and_log(output)
def change_index_file_permission(public_ip, key_path):
ssh_permission_cmd = construct_ssh(
key_path, public_ip, " 'sudo chmod 646 /usr/share/nginx/html/index.html'")
utils.print_and_log('Trying to change permission on index.html')
(status, output) = subprocess.getstatusoutput(ssh_permission_cmd)
if status == 0:
utils.print_and_log('Successfully changed the file permission')
else:
utils.print_and_log('Failed to change file permission')
utils.print_and_log(output)
def get_collection(coll_name, db, projection=None, query_filter=None, limit=None):
"""
Downloads and returns single collection from MongoDB and returns dataframe.
Optional query filter can be applied (useful for downloading logins post-views from events table.
Returns a dataframe.
"""
kwargs = {} # necessary because pymongo function doesn't accept limit=None
if limit:
kwargs.update({'limit':limit})
print_and_log('{} download started at {}. . .'.format(coll_name, datetime.datetime.today()))
coll_list = db[coll_name].find(query_filter, projection=projection, **kwargs)
coll_df = pd.DataFrame(list(coll_list))
print_and_log('{} download completed at {}!'.format(coll_name, datetime.datetime.today()))
return coll_df
def __init__(self):
self.args = self.parse_command_line()
self.checkpoint_dir, self.logfile, self.checkpoint_path_validation, self.checkpoint_path_final \
= get_log_files(self.args.checkpoint_dir, self.args.resume_from_checkpoint, False)
print_and_log(self.logfile, "Options: %s\n" % self.args)
print_and_log(self.logfile,
"Checkpoint Directory: %s\n" % self.checkpoint_dir)
self.writer = SummaryWriter()
#gpu_device = 'cuda:0'
gpu_device = 'cuda'
self.device = torch.device(
gpu_device if torch.cuda.is_available() else 'cpu')
self.model = self.init_model()
self.train_set, self.validation_set, self.test_set = self.init_data()
self.vd = video_reader.VideoDataset(self.args)
self.video_loader = torch.utils.data.DataLoader(
self.vd, batch_size=1, num_workers=self.args.num_workers)
self.loss = loss
self.accuracy_fn = aggregate_accuracy
if self.args.opt == "adam":
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.args.learning_rate)
elif self.args.opt == "sgd":
self.optimizer = torch.optim.SGD(self.model.parameters(),
lr=self.args.learning_rate)
self.test_accuracies = TestAccuracies(self.test_set)
self.scheduler = MultiStepLR(self.optimizer,
milestones=self.args.sch,
gamma=0.1)
self.start_iteration = 0
if self.args.resume_from_checkpoint:
self.load_checkpoint()
self.optimizer.zero_grad()
def init_gan_model(model_dir, nb_of_classes):
device = torch.device("cuda:0")
generator_net = networks.Generator().to(device)
discriminator_net = networks.ClassifierDiscriminator(nb_of_classes=nb_of_classes).to(device)
generator_net.apply(networks.weights_init)
discriminator_net.apply(networks.weights_init)
utils.print_and_log(model_dir, generator_net)
utils.print_and_log(model_dir, discriminator_net)
learning_rate = 0.0002
beta1 = 0.5
discriminator_optimizer = optim.Adam(discriminator_net.parameters(), lr=learning_rate, betas=(beta1, 0.999))
generator_optimizer = optim.Adam(generator_net.parameters(), lr=learning_rate, betas=(beta1, 0.999))
criterion = nn.CrossEntropyLoss()
return discriminator_net, generator_net, criterion, discriminator_optimizer, generator_optimizer
def __init__(self):
self.args = self.parse_command_line()
self.checkpoint_dir, self.logfile, self.checkpoint_path_validation, self.checkpoint_path_final \
= get_log_files(self.args.checkpoint_dir)
print_and_log(self.logfile, "Options: %s\n" % self.args)
print_and_log(self.logfile, "Checkpoint Directory: %s\n" % self.checkpoint_dir)
gpu_device = 'cuda:0'
self.device = torch.device(gpu_device if torch.cuda.is_available() else 'cpu')
self.model = self.init_model()
self.train_set, self.validation_set, self.test_set = self.init_data()
self.metadataset = MetaDatasetReader(self.args.data_path, self.args.mode, self.train_set, self.validation_set,
self.test_set)
self.loss = loss
self.accuracy_fn = aggregate_accuracy
self.optimizer = torch.optim.Adam(self.model.parameters(), lr=self.args.learning_rate)
self.optimizer.zero_grad()
self.validation_accuracies = ValidationAccuracies(self.validation_set)
def start_service(service_dir, branch, container_name, flag_str, log=None):
log= print_and_log("[*] Starting service from %s (branch '%s')" % \
(service_dir, branch), log)
checkout(service_dir, branch)
# Update flag file
flag_path = os.path.join(service_dir, "flag") # Assumption in template
if not os.path.isfile(flag_path):
log = print_and_log("[*] 'flag' file not found in %s" % service_dir,
log)
return False, log
with open(flag_path, "w") as flag_file:
flag_file.write(flag_str)
# Run the service
script = os.path.join(base_dir(), "setup_service.sh")
cmdline = \
"%s %s %d %d" % (script, container_name, SERVICE_PORT, SERVICE_PORT)
output, err, e = run_command(cmdline, service_dir)
if e != 0:
log = print_and_log("[*] Failed to start service", log)
log = print_and_log(err, log)
log = print_and_log("==========================", log)
return False, log
if log is not None:
log = log + output
log = print_and_log("[*] Started service successfully", log)
return True, log
def run_dash_aggregations_pipeline(collections, date_str):
plot_specs = generate_specs(collections)
plot_specs_small = [
PlotSpec(
title=spec.title,
data=pd.DataFrame(),
date_column=spec.date_column,
agg_func=spec.agg_func,
agg_column=spec.agg_column,
color=spec.color,
) for spec in plot_specs
]
timeseries_dict = generate_timeseries_dict(plot_specs)
directory = BASE_PATH + '{folder}/{date}'.format(folder='processed',
date=date_str)
print_and_log('Writing timeseries_dict to disk.')
pickle.dump(plot_specs_small,
open(directory + '/plot_specs.p', 'wb'),
protocol=pickle.HIGHEST_PROTOCOL)
print_and_log('Writing plot_specs to disk.')
pickle.dump(timeseries_dict,
open(directory + '/timeseries_dict.p', 'wb'),
protocol=pickle.HIGHEST_PROTOCOL)
print_and_log('Writing timeseries_dict and plot_specs to disk completed')
def test(self, path, session):
print_and_log(self.logfile, "") # add a blank line
print_and_log(self.logfile, 'Testing model {0:}: '.format(path))
self.model = self.init_model()
self.model.load_state_dict(torch.load(path))
with torch.no_grad():
for item in self.test_set:
accuracies = []
for _ in range(NUM_TEST_TASKS):
task_dict = self.dataset.get_test_task(item, session)
context_images, target_images, context_labels, target_labels = self.prepare_task(
task_dict)
target_logits = self.model(context_images, context_labels,
target_images)
accuracy = self.accuracy_fn(target_logits, target_labels)
accuracies.append(accuracy.item())
del target_logits
accuracy = np.array(accuracies).mean() * 100.0
accuracy_confidence = (196.0 *
np.array(accuracies).std()) / np.sqrt(
len(accuracies))
print_and_log(
self.logfile,
'{0:}: {1:3.1f}+/-{2:2.1f}'.format(item, accuracy,
accuracy_confidence))
def run_exploit(exploit_dir, container_name, timeout, log=None):
log = print_and_log("[*] Running exploit", log)
script = os.path.join(base_dir(), "launch_exploit.sh")
cmdline = \
"%s %s %s %d %d" % \
(script, container_name, SERVICE_IP, SERVICE_PORT, timeout)
output, err, e = run_command(cmdline, exploit_dir)
if log is not None:
log = log + output
if e != 0:
log = print_and_log("[*] Failed to run exploit", log)
log = print_and_log(err, log)
log = print_and_log("==========================", log)
return None, log
# Exploit prints out the flag string at the end.
tokens = [_f for _f in output.split('\n')
if _f] # Filter out empty strings
flag_candidate = [_f for _f in tokens if _f][-1] # Read the last line
return flag_candidate, log
def init_model(model_dir):
device = torch.device("cuda:0")
generator_net = networks.Generator().to(device)
discriminator_net = networks.Discriminator().to(device)
generator_net.apply(networks.weights_init)
discriminator_net.apply(networks.weights_init)
utils.print_and_log(model_dir, generator_net)
utils.print_and_log(model_dir, discriminator_net)
learning_rate = 0.0002
beta1 = 0.5
discriminator_optimizer = optim.Adam(discriminator_net.parameters(),
lr=learning_rate,
betas=(beta1, 0.999))
generator_optimizer = optim.Adam(generator_net.parameters(),
lr=learning_rate,
betas=(beta1, 0.999))
return discriminator_net, generator_net, discriminator_optimizer, generator_optimizer
def bulk_upload_to_pg(df, table_name, conn=None, clean_text=True):
print_and_log('Upload {}'.format(table_name))
df = df.copy()
df.loc[:, 'birth'] = pd.datetime.now()
df.columns = df.columns.to_series().apply(camel_to_snake)
if clean_text:
df = clean_dataframe_text(df)
sep = '\t'
buffer = StringIO()
buffer.write(
df.to_csv(index=None,
header=None,
sep=sep,
na_rep='',
escapechar='\\',
quoting=csv.QUOTE_NONE)
) # Write the Pandas DataFrame as a csv to the buffer
buffer.seek(0) # Be sure to reset the position to the start of the stream
def execute_copy(conn):
dbapi_conn = conn.connection
with dbapi_conn.cursor() as c:
c.copy_from(buffer,
table_name,
columns=df.columns,
sep=sep,
null='')
if not conn:
engine = get_pg_engine()
with engine.begin() as conn:
execute_copy(conn)
else:
execute_copy(conn)
def __init__(self):
self.args = self.parse_command_line()
self.checkpoint_dir, self.logfile, self.checkpoint_path_validation, self.checkpoint_path_final \
= get_log_files(self.args.checkpoint_dir, self.args.resume_from_checkpoint, self.args.mode == "test" or
self.args.mode == "attack")
print_and_log(self.logfile, "Options: %s\n" % self.args)
print_and_log(self.logfile,
"Checkpoint Directory: %s\n" % self.checkpoint_dir)
gpu_device = 'cuda:0'
self.device = torch.device(
gpu_device if torch.cuda.is_available() else 'cpu')
self.model = self.init_model()
self.train_set, self.validation_set, self.test_set = self.init_data()
if self.args.dataset == "meta-dataset":
self.dataset = MetaDatasetReader(
self.args.data_path, self.args.mode, self.train_set,
self.validation_set, self.test_set, self.args.max_way_train,
self.args.max_way_test, self.args.max_support_train,
self.args.max_support_test)
else:
self.dataset = SingleDatasetReader(self.args.data_path,
self.args.mode,
self.args.dataset,
self.args.way, self.args.shot,
self.args.query_train,
self.args.query_test)
self.loss = loss
self.accuracy_fn = aggregate_accuracy
self.optimizer = torch.optim.Adam(self.model.parameters(),
lr=self.args.learning_rate)
self.validation_accuracies = ValidationAccuracies(self.validation_set)
self.start_iteration = 0
if self.args.resume_from_checkpoint:
self.load_checkpoint()
self.optimizer.zero_grad()
def main():
while True:
menu()
choice = input("\nEnter your choice: ")
if choice == "1":
print("Create instance and bucket")
create_instance()
create_bucket()
elif choice == "2":
print("Create instance")
create_instance()
elif choice == "3":
print("Create bucket")
create_bucket()
elif choice == "4":
print("Upload to bucket")
list_and_upload_bucket()
elif choice == "5":
try:
public_ip = get_instance_ip()
if public_ip:
key_path = make_key_read_only(utils.get_valid_key("Enter path to your private key: "))
run_check_webserver(public_ip, key_path)
except Exception as error:
utils.print_and_log("Run Check Server Error: " + str(error))
elif choice == '6':
try:
public_ip = get_instance_ip()
if public_ip:
key_path = make_key_read_only(utils.get_valid_key("Enter path to your private key: "))
get_instance_usage(key_path, public_ip)
except Exception as error:
utils.print_and_log("Get Instance Usage Error: " + str(error))
elif choice == "0":
print("Exiting")
sys.exit(0)
else:
print("Not a valid choice")
def run(self):
if self.args.mode == 'train' or self.args.mode == 'train_test':
train_accuracies = []
losses = []
total_iterations = self.args.training_iterations
for iteration in range(self.start_iteration, total_iterations):
torch.set_grad_enabled(True)
task_dict = self.dataset.get_train_task()
task_loss, task_accuracy = self.train_task(task_dict)
train_accuracies.append(task_accuracy)
losses.append(task_loss)
# optimize
if ((iteration + 1) % self.args.tasks_per_batch
== 0) or (iteration == (total_iterations - 1)):
self.optimizer.step()
self.optimizer.zero_grad()
if (iteration + 1) % PRINT_FREQUENCY == 0:
# print training stats
print_and_log(
self.logfile,
'Task [{}/{}], Train Loss: {:.7f}, Train Accuracy: {:.7f}'
.format(iteration + 1, total_iterations,
torch.Tensor(losses).mean().item(),
torch.Tensor(train_accuracies).mean().item()))
train_accuracies = []
losses = []
if ((iteration + 1) % self.args.val_freq
== 0) and (iteration + 1) != total_iterations:
# validate
accuracy_dict = self.validate()
self.validation_accuracies.print(self.logfile,
accuracy_dict)
# save the model if validation is the best so far
if self.validation_accuracies.is_better(accuracy_dict):
self.validation_accuracies.replace(accuracy_dict)
torch.save(self.model.state_dict(),
self.checkpoint_path_validation)
print_and_log(self.logfile,
'Best validation model was updated.')
print_and_log(self.logfile, '')
self.save_checkpoint(iteration + 1)
# save the final model
torch.save(self.model.state_dict(), self.checkpoint_path_final)
if self.args.mode == 'train_test':
self.test(self.checkpoint_path_final)
self.test(self.checkpoint_path_validation)
if self.args.mode == 'test':
self.test(self.args.test_model_path)
self.logfile.close()
