def train_network(network, device):
csv_reader = CSVReader()
# Load the training data from the csv files
training_loader = csv_reader.get_training_loader(
Config.batch_size, Config.shuffle, Config.load_data_from_file,
Config.preprocess_data)
validation_loader = csv_reader.get_validation_loader(
Config.batch_size, Config.shuffle, Config.load_data_from_file,
Config.preprocess_data)
# Converts data into 3D renderings
if (Config.use_image_generator):
from skeleton_visualizer import SkeletonVisualizer
# You need to give the visualizer 2 skeleton instances to display
skeleton1 = training_loader.dataset.features[0 + 12].numpy()
skeleton2 = training_loader.dataset.features[6261 + 12].numpy()
render = SkeletonVisualizer(skeleton1, skeleton2)
return
optimizer = optim.SGD(net.parameters(), lr=0.01)
criterion_ce = nn.CrossEntropyLoss(size_average=True)
num_classes = RecognitionDataset.num_classifications
class_names = RecognitionDataset.classification_names
# This plotter is a class than I have created
visdom_env_name = 'Plots'
if (Config.use_visdom):
plotter = visdom_utils.VisdomLinePlotter(env_name=visdom_env_name)
confusion_logger_training = torchnet.logger.VisdomLogger(
'heatmap',
env=visdom_env_name,
opts={
'title': 'Training Confusion Matrix',
'columnnames': class_names,
'rownames': class_names
})
confusion_logger_validation = torchnet.logger.VisdomLogger(
'heatmap',
env=visdom_env_name,
opts={
'title': 'Validation Confusion Matrix',
'columnnames': class_names,
'rownames': class_names
})
confusion_graph_training = torchnet.meter.ConfusionMeter(num_classes)
confusion_graph_validation = torchnet.meter.ConfusionMeter(num_classes)
batch_count = 1
total_confusion = np.zeros((num_classes, num_classes))
for epoch in range(Config.max_num_epochs):
total = 0
total_correct = 0
network.train()
if (Config.use_visdom):
confusion_graph_training.reset()
confusion_graph_validation.reset()
for data in training_loader:
features, labels = data
#print((features.min(), features.max()))
loss_ce, target_np, predicted_np = run_network(
network, optimizer, criterion_ce, device, features, labels,
batch_count, True)
confusion = np.swapaxes(
confusion_matrix(target_np,
predicted_np,
labels=np.arange(0, num_classes)), 0, 1)
total_confusion = np.add(total_confusion, confusion)
# The predicted and target are flopped in order to make the data on the correct axis
#confusion_graph_training.add(torch.from_numpy(target_np), torch.from_numpy(predicted_np))
if (Config.use_visdom):
confusion_graph_training.add(torch.from_numpy(predicted_np),
torch.from_numpy(target_np))
total_correct += (predicted_np == target_np).sum()
total += features.size(0)
batch_count += 1
# Calculate the network's accuracy every epoch
validation_correct, validation_total = get_validation_accuracy(
network, optimizer, criterion_ce, device, validation_loader,
confusion_graph_validation)
accuracy_validation = 100 * validation_correct / validation_total
accuracy_training = 100 * total_correct / total
print(
'loss={} TRAINING: #total={} #correct={} VALIDATION: #total={} #correct={}'
.format(loss_ce, total, total_correct, validation_total,
validation_correct))
# Plot our results to visdom
if (Config.use_visdom):
plotter.plot('loss', 'train', 'Class Loss', epoch, loss_ce.item())
plotter.plot('accuracy', 'train', 'Class Accuracy', epoch,
accuracy_training)
plotter.plot('accuracy', 'validation', 'Class Accuracy', epoch,
accuracy_validation)
confusion_logger_training.log(confusion_graph_training.value())
confusion_logger_validation.log(confusion_graph_validation.value())
