159 lines
4.4 KiB
C++
159 lines
4.4 KiB
C++
#include <getopt.h>
|
|
#include <fstream>
|
|
#include <memory>
|
|
#include <fcntl.h>
|
|
#include <google/protobuf/text_format.h>
|
|
#include <google/protobuf/io/zero_copy_stream_impl.h>
|
|
#include <caffe/proto/caffe.pb.h>
|
|
|
|
static struct
|
|
{
|
|
bool verbose = false;
|
|
std::string_view output_filename;
|
|
std::string_view input_filename;
|
|
} options;
|
|
|
|
template<typename... Args>
|
|
void verbose_print(Args &&... args)
|
|
{
|
|
if (!options.verbose) return;
|
|
|
|
(std::cerr << ... << args) << '\n';
|
|
}
|
|
|
|
void show_help(std::string_view name, int exit_code)
|
|
{
|
|
std::cerr << "Usage: " << name << "[OPTIONS] [FILENAME]\n"
|
|
"\n"
|
|
"Valid options:\n"
|
|
"-h\tshow this message\n"
|
|
"-v\tshow debug messages\n"
|
|
"-o OUTPUT\t write output to OUTPUT\n"
|
|
"FILENAME\t read from FILENAME instead of stdin\n" << std::endl;
|
|
|
|
exit(exit_code);
|
|
}
|
|
|
|
int get_file_descriptor(std::string_view file_name, int mode)
|
|
{
|
|
int fd = open(file_name.data(), mode);
|
|
if (fd < 0) {
|
|
perror("get_file_descriptor");
|
|
exit(1);
|
|
}
|
|
|
|
return fd;
|
|
}
|
|
|
|
void read_options(int argc, char **argv)
|
|
{
|
|
for (char c; (c = static_cast<char>(getopt(argc, argv, "hvo:"))) != -1;) {
|
|
switch (c) {
|
|
case 'v':
|
|
options.verbose = true;
|
|
break;
|
|
case 'h':
|
|
show_help(argv[0], 0);
|
|
break;
|
|
case 'o':
|
|
options.output_filename = optarg;
|
|
break;
|
|
case '?':
|
|
show_help(argv[0], 1);
|
|
break;
|
|
}
|
|
}
|
|
|
|
if (argc > optind) {
|
|
options.input_filename = argv[optind];
|
|
}
|
|
}
|
|
|
|
void read_network(caffe::NetParameter &network)
|
|
{
|
|
using namespace google::protobuf::io;
|
|
std::unique_ptr<ZeroCopyInputStream> input;
|
|
if (!options.input_filename.empty()) {
|
|
verbose_print("Reading network from", options.input_filename);
|
|
|
|
auto fd = get_file_descriptor(options.input_filename, O_RDONLY);
|
|
FileInputStream *i = new FileInputStream(fd);
|
|
i->SetCloseOnDelete(true);
|
|
input.reset(i);
|
|
} else {
|
|
verbose_print("Reading network from stdin");
|
|
input = std::make_unique<IstreamInputStream>(&std::cin);
|
|
}
|
|
|
|
auto result = google::protobuf::TextFormat::Parse(input.get(), &network);
|
|
if (!result) {
|
|
std::cerr << "Error reading network file!" << std::endl;
|
|
exit(2);
|
|
}
|
|
}
|
|
|
|
template<typename LayerType>
|
|
void patch_layers(google::protobuf::RepeatedPtrField<LayerType> *layers)
|
|
{
|
|
std::map<std::string, std::string> outputs;
|
|
|
|
for (LayerType &layer : *layers) {
|
|
|
|
for (int i = 0; i < layer.bottom_size(); ++i) {
|
|
if (auto it = outputs.find(layer.bottom(i)); it != outputs.end()) {
|
|
verbose_print(layer.name(), " reads from in-place ", layer.bottom(i), ", rewriting.");
|
|
*layer.mutable_bottom(i) = it->second;
|
|
}
|
|
}
|
|
|
|
for (int i = 0; i < layer.top_size(); ++i) {
|
|
if (auto it = outputs.find(layer.top(i)); it != outputs.end()) {
|
|
verbose_print(layer.name(), " works in-place rewriting.");
|
|
it->second = layer.name();
|
|
*layer.mutable_top(i) = layer.name();
|
|
} else {
|
|
outputs[layer.name()] = layer.name();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
void patch_network(caffe::NetParameter &network)
|
|
{
|
|
patch_layers(network.mutable_layer());
|
|
patch_layers(network.mutable_layers());
|
|
}
|
|
|
|
void write_network(caffe::NetParameter &network)
|
|
{
|
|
using namespace google::protobuf::io;
|
|
std::unique_ptr<ZeroCopyOutputStream> output;
|
|
if (!options.output_filename.empty()) {
|
|
verbose_print("Writing network to ", options.output_filename);
|
|
|
|
auto fd = get_file_descriptor(options.output_filename, O_WRONLY | O_CREAT | O_TRUNC);
|
|
FileOutputStream *i = new FileOutputStream(fd);
|
|
i->SetCloseOnDelete(true);
|
|
output.reset(i);
|
|
} else {
|
|
verbose_print("Writing network to stdout");
|
|
output = std::make_unique<OstreamOutputStream>(&std::cout);
|
|
}
|
|
|
|
google::protobuf::TextFormat::Print(network, output.get());
|
|
}
|
|
|
|
int main(int argc, char **argv)
|
|
{
|
|
GOOGLE_PROTOBUF_VERIFY_VERSION;
|
|
|
|
read_options(argc, argv);
|
|
caffe::NetParameter network;
|
|
read_network(network);
|
|
patch_network(network);
|
|
write_network(network);
|
|
|
|
// Deallocate all global protobuf objects.
|
|
google::protobuf::ShutdownProtobufLibrary();
|
|
}
|