def train(load_model: str, save_model: str, train_dataset: str,
test_dataset: str, no_train: bool, no_test: bool, epochs: int,
batch_size: int, learning_rate: float):
device = torch.device('cuda:0' if cuda.is_available() else 'cpu')
click.secho('Using device={}'.format(device), fg='blue')
net = Net()
net.to(device)
if load_model is not None:
click.secho('Loading model from \'{}\''.format(load_model),
fg='yellow')
net.load_state_dict(torch.load(load_model, map_location=device))
if not no_train:
click.echo('Training model using {}'.format(train_dataset))
net.train()
train_net(net,
data_path=train_dataset,
batch_size=batch_size,
num_epochs=epochs,
learning_rate=learning_rate)
if not no_train and save_model is not None:
click.secho('Saving model as \'{}\''.format(save_model), fg='yellow')
torch.save(net.state_dict(), save_model)
if not no_test:
click.echo('Testing model using {}'.format(test_dataset))
net.eval()
accuracy = test_net(net, data_path=test_dataset, batch_size=batch_size)
color = 'green' if accuracy > 97. else 'red'
click.secho('Accuracy={}'.format(accuracy), fg=color)
def visualize(model: str, images: [str], occlusion_window: int,
occlusion_stride: int, no_occlussion: bool, no_gradient: bool):
device = torch.device('cuda:0' if cuda.is_available() else 'cpu')
click.secho('Using device={}'.format(device), fg='blue')
net = Net()
net.to(device)
click.secho('Loading model from \'{}\''.format(model), fg='yellow')
net.load_state_dict(torch.load(model, map_location=device))
net.eval()
for path in images:
image = utils.load_image(path).to(device)
output = net(image)
_, predicted = torch.max(output.data, 1)
click.echo('Image \'{}\' most likely represents a \'{}\''.format(
path, classes[predicted]))
if not no_occlussion:
occlustion(net,
image,
predicted,
k=occlusion_window,
stride=occlusion_stride)
if not no_gradient:
gradient(net, image, predicted)
def main():
"""Main function
"""
# Load the parameters
args = args_parser()
json_path = os.path.join(args.model_dir, 'params.json')
assert os.path.isfile(
json_path), "No json configuration file found at {}".format(json_path)
params = utils.Params(json_path)
# Create summary writer for use with tensorboard
writer = SummaryWriter(os.path.join(args.model_dir, 'runs', 'eval'))
# use GPU if available
params.cuda = torch.cuda.is_available() # use GPU is available
# Set the random seed for reproducible experiments
torch.manual_seed(230)
if params.cuda:
torch.cuda.manual_seed(230)
params.device = "cuda:0"
else:
params.device = "cpu"
# Set the logger
utils.set_logger(os.path.join(args.model_dir, 'evaluate.log'))
logging.info("Loading the dataset...")
# fetch dataloaders
dataloaders = d_l.get_dataloader(['test'], args.data_dir, params)
test_dl = dataloaders['test']
logging.info("- done.")
# Define the model
model = Net(params)
if params.cuda:
model = model.to(params.device)
writer.add_graph(model, next(iter(test_dl))[0])
criterion = loss_fn
metrics = get_metrics()
logging.info("Starting evaluation")
# Reload weights from the saved file
utils.load_checkpoint(
os.path.join(args.model_dir, args.restore_file + '.pth.tar'), model)
# Evaluate
test_metrics = evaluate(model, criterion, test_dl, metrics, params, writer,
0)
save_path = os.path.join(args.model_dir,
"metrics_test_{}.json".format(args.restore_file))
utils.save_dict_to_json(test_metrics, save_path)
writer.close()
def main():
"""Main function
"""
# Load the parameters from json file
args = args_parser()
json_path = os.path.join(args.model_dir, 'params.json')
assert os.path.isfile(
json_path), "No json configuration file found at {}".format(json_path)
params = utils.Params(json_path)
# Create summary writer for use with tensorboard
writer = SummaryWriter(os.path.join(args.model_dir, 'runs', 'train'))
# use GPU if available
params.cuda = torch.cuda.is_available()
# Set the random seed for reproducible experiments
torch.manual_seed(230)
if params.cuda:
torch.cuda.manual_seed(230)
params.device = "cuda:0"
else:
params.device = "cpu"
# Set the logger
utils.set_logger(os.path.join(args.model_dir, 'train.log'))
# Create the input data pipeline
logging.info("Loading the datasets...")
# fetch dataloaders
dataloaders = d_l.get_dataloader(['train', 'val'], args.data_dir, params)
train_dl = dataloaders['train']
val_dl = dataloaders['val']
logging.info("- done.")
# Define the model and optimizer
model = Net(params)
if params.cuda:
model = model.to(params.device)
writer.add_graph(model, next(iter(train_dl))[0])
optimizer = torch.optim.Adam(model.parameters(), lr=params.learning_rate)
# fetch loss function and metrics
criterion = loss_fn
metrics = get_metrics()
# Train the model
logging.info("Starting training for %d epoch(s)", params.num_epochs)
train_and_evaluate(model, train_dl, val_dl, optimizer, criterion, metrics,
params, args.model_dir, writer, args.restore_file)
writer.close()
