Morph MORPH
SIGN IN SIGN UP
pytorch / examples UNCLAIMED

A set of examples around pytorch in Vision, Text, Reinforcement Learning, etc.

0 0 0 Python
Normal View History Raw
Added Transfer Learning using Libtorch and OpenCV (#626) * Transfer Learning Example using Libtorch and OpenCV * Revised CMakeLists file * Updated README * Removed copyrights * Removed copyrights * Removed copyrights * Update classify.cpp * Update main.cpp * Update main.cpp * Update classify.cpp
2019-09-03 22:03:15 +01:00
//
// main.cpp
// transfer-learning
//
// Created by Kushashwa Ravi Shrimali on 12/08/19.
//
#include "main.h"
torch::Tensor read_data(std::string location) {
/*
Function to return image read at location given as type torch::Tensor
Resizes image to (224, 224, 3)
Parameters
===========
1. location (std::string type) - required to load image from the location
Returns
===========
torch::Tensor type - image read as tensor
*/
cv::Mat img = cv::imread(location, 1);
cv::resize(img, img, cv::Size(224, 224), cv::INTER_CUBIC);
torch::Tensor img_tensor = torch::from_blob(img.data, {img.rows, img.cols, 3}, torch::kByte);
img_tensor = img_tensor.permute({2, 0, 1});
return img_tensor.clone();
}
torch::Tensor read_label(int label) {
/*
Function to return label from int (0, 1 for binary and 0, 1, ..., n-1 for n-class classification) as type torch::Tensor
Parameters
===========
1. label (int type) - required to convert int to tensor
Returns
===========
torch::Tensor type - label read as tensor
*/
torch::Tensor label_tensor = torch::full({1}, label);
return label_tensor.clone();
}
std::vector<torch::Tensor> process_images(std::vector<std::string> list_images) {
/*
Function returns vector of tensors (images) read from the list of images in a folder
Parameters
===========
1. list_images (std::vector<std::string> type) - list of image paths in a folder to be read
Returns
===========
std::vector<torch::Tensor> type - Images read as tensors
*/
std::vector<torch::Tensor> states;
for(std::vector<std::string>::iterator it = list_images.begin(); it != list_images.end(); ++it) {
torch::Tensor img = read_data(*it);
states.push_back(img);
}
return states;
}
std::vector<torch::Tensor> process_labels(std::vector<int> list_labels) {
/*
Function returns vector of tensors (labels) read from the list of labels
Parameters
===========
1. list_labels (std::vector<int> list_labels) -
Returns
===========
std::vector<torch::Tensor> type - returns vector of tensors (labels)
*/
std::vector<torch::Tensor> labels;
for(std::vector<int>::iterator it = list_labels.begin(); it != list_labels.end(); ++it) {
torch::Tensor label = read_label(*it);
labels.push_back(label);
}
return labels;
}
std::pair<std::vector<std::string>,std::vector<int>> load_data_from_folder(std::vector<std::string> folders_name) {
/*
Function to load data from given folder(s) name(s) (folders_name)
Returns pair of vectors of string (image locations) and int (respective labels)
Parameters
===========
1. folders_name (std::vector<std::string> type) - name of folders as a vector to load data from
Returns
===========
std::pair<std::vector<std::string>, std::vector<int>> type - returns pair of vector of strings (image paths) and respective labels' vector (int label)
*/
std::vector<std::string> list_images;
std::vector<int> list_labels;
int label = 0;
for(auto const& value: folders_name) {
std::string base_name = value + "/";
// cout << "Reading from: " << base_name << endl;
DIR* dir;
struct dirent *ent;
if((dir = opendir(base_name.c_str())) != NULL) {
while((ent = readdir(dir)) != NULL) {
std::string filename = ent->d_name;
if(filename.length() > 4 && filename.substr(filename.length() - 3) == "jpg") {
// cout << base_name + ent->d_name << endl;
// cv::Mat temp = cv::imread(base_name + "/" + ent->d_name, 1);
list_images.push_back(base_name + ent->d_name);
list_labels.push_back(label);
}
}
closedir(dir);
} else {
std::cout << "Could not open directory" << std::endl;
// return EXIT_FAILURE;
}
label += 1;
}
return std::make_pair(list_images, list_labels);
}
template<typename Dataloader>
void train(torch::jit::script::Module net, torch::nn::Linear lin, Dataloader& data_loader, torch::optim::Optimizer& optimizer, size_t dataset_size) {
/*
This function trains the network on our data loader using optimizer.
Also saves the model as model.pt after every epoch.
Parameters
===========
1. net (torch::jit::script::Module type) - Pre-trained model without last FC layer
2. lin (torch::nn::Linear type) - last FC layer with revised out_features depending on the number of classes
3. data_loader (DataLoader& type) - Training data loader
4. optimizer (torch::optim::Optimizer& type) - Optimizer like Adam, SGD etc.
5. size_t (dataset_size type) - Size of training dataset
Returns
===========
Nothing (void)
*/
float best_accuracy = 0.0;
int batch_index = 0;
for(int i=0; i<25; i++) {
float mse = 0;
float Acc = 0.0;
for(auto& batch: *data_loader) {
auto data = batch.data;
auto target = batch.target.squeeze();
// Should be of length: batch_size
data = data.to(torch::kF32);
target = target.to(torch::kInt64);
std::vector<torch::jit::IValue> input;
input.push_back(data);
optimizer.zero_grad();
auto output = net.forward(input).toTensor();
// For transfer learning
output = output.view({output.size(0), -1});
output = lin(output);
auto loss = torch::nll_loss(torch::log_softmax(output, 1), target);
loss.backward();
optimizer.step();
auto acc = output.argmax(1).eq(target).sum();
Acc += acc.template item<float>();
mse += loss.template item<float>();
batch_index += 1;
}
mse = mse/float(batch_index); // Take mean of loss
std::cout << "Epoch: " << i << ", " << "Accuracy: " << Acc/dataset_size << ", " << "MSE: " << mse << std::endl;
test(net, lin, data_loader, dataset_size);
if(Acc/dataset_size > best_accuracy) {
best_accuracy = Acc/dataset_size;
std::cout << "Saving model" << std::endl;
net.save("model.pt");
torch::save(lin, "model_linear.pt");
}
}
}
template<typename Dataloader>
void test(torch::jit::script::Module network, torch::nn::Linear lin, Dataloader& loader, size_t data_size) {
/*
Function to test the network on test data
Parameters
===========
1. network (torch::jit::script::Module type) - Pre-trained model without last FC layer
2. lin (torch::nn::Linear type) - last FC layer with revised out_features depending on the number of classes
3. loader (Dataloader& type) - test data loader
4. data_size (size_t type) - test data size
Returns
===========
Nothing (void)
*/
network.eval();
float Loss = 0, Acc = 0;
for (const auto& batch : *loader) {
auto data = batch.data;
auto targets = batch.target.squeeze();
data = data.to(torch::kF32);
targets = targets.to(torch::kInt64);
std::vector<torch::jit::IValue> input;
input.push_back(data);
auto output = network.forward(input).toTensor();
output = output.view({output.size(0), -1});
output = lin(output);
auto loss = torch::nll_loss(torch::log_softmax(output, 1), targets);
auto acc = output.argmax(1).eq(targets).sum();
Loss += loss.template item<float>();
Acc += acc.template item<float>();
}
std::cout << "Test Loss: " << Loss/data_size << ", Acc:" << Acc/data_size << std::endl;
}
int main(int argc, const char * argv[]) {
// Set folder names for cat and dog images
std::string cats_name = "/Users/krshrimali/Documents/krshrimali-blogs/dataset/train/cat_test";
std::string dogs_name = "/Users/krshrimali/Documents/krshrimali-blogs/dataset/train/dog_test";
std::vector<std::string> folders_name;
folders_name.push_back(cats_name);
folders_name.push_back(dogs_name);
// Get paths of images and labels as int from the folder paths
std::pair<std::vector<std::string>, std::vector<int>> pair_images_labels = load_data_from_folder(folders_name);
std::vector<std::string> list_images = pair_images_labels.first;
std::vector<int> list_labels = pair_images_labels.second;
// Initialize CustomDataset class and read data
auto custom_dataset = CustomDataset(list_images, list_labels).map(torch::data::transforms::Stack<>());
// Load pre-trained model
// You can also use: auto module = torch::jit::load(argv[1]);
torch::jit::script::Module module = torch::jit::load(argv[1]);
// Resource: https://discuss.pytorch.org/t/how-to-load-the-prebuilt-resnet-models-or-any-other-prebuilt-models/40269/8
// For VGG: 512 * 14 * 14, 2
torch::nn::Linear lin(512, 2); // the last layer of resnet, which we want to replace, has dimensions 512x1000
torch::optim::Adam opt(lin->parameters(), torch::optim::AdamOptions(1e-3 /*learning rate*/));
auto data_loader = torch::data::make_data_loader<torch::data::samplers::RandomSampler>(std::move(custom_dataset), 4);
train(module, lin, data_loader, opt, custom_dataset.size().value());
return 0;
}
Reference in New Issue Copy Permalink