Pytorch c++ cmake file and train model

A Minimal CmakeFile.txt ( with libtorch C++ ): find_package

cmake_minimum_required(VERSION 3.14 FATAL_ERROR)

project(linear-regression VERSION 1.0.0 LANGUAGES CXX)

find_package(Torch REQUIRED)

set(EXECUTABLE_NAME linear-regression)

add_executable(${EXECUTABLE_NAME})
target_sources(${EXECUTABLE_NAME} PRIVATE main.cpp)

target_link_libraries(${EXECUTABLE_NAME} ${TORCH_LIBRARIES})

set_target_properties(${EXECUTABLE_NAME} PROPERTIES
  CXX_STANDARD 14
  CXX_STANDARD_REQUIRED YES
)

if(MSVC)
    set_target_properties(${EXECUTABLE_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY_RELEASE ${CMAKE_CURRENT_BINARY_DIR})
    set_target_properties(${EXECUTABLE_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY_DEBUG ${CMAKE_CURRENT_BINARY_DIR})
    include(copy_torch_dlls)
    copy_torch_dlls(${EXECUTABLE_NAME})
endif(MSVC)

Another ```CMakeList.txt`` file for single .cpp file (for example compile example.cpp)

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(example)
SET(CMAKE_BUILD_TYPE Release)
set(EXECUTABLE_OUTPUT_PATH ${
   PROJECT_SOURCE_DIR}/bin)
set(LIBRARY_OUTPUT_PATH ${
   PROJECT_SOURCE_DIR}/lib)

find_package(OpenCV REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
set(CMAKE_PREFIX_PATH /home/jerry/libtorch)
find_package(Torch REQUIRED)
add_executable(example src/example.cpp)

include_directories(
${
   PROJECT_SOURCE_DIR}/include
${
   OpenCV_INCLUDE_DIRS}
)

target_link_libraries(example ${
   OpenCV_LIBS} ${
   TORCH_LIBRARIES})
set_property(TARGET example PROPERTY CXX_STANDARD 14)

Remember to use find_package(Torch REQUIRED)
And also: set export Torch_dir=/home/jerry/libtorch/share/cmake/Torch(because config.cmake is at this directory)

A minimal Makefile :

remember to add a objs folder and the following will output executable ./start

CXX=g++

EIGEN_LOCATION=/usr/include/eigen3

BUILD_DIR=objs

CXXFLAGS =-Wall
CXXFLAGS+=-O3
CXXFLAGS+=-std=c++11
CXXFLAGS+=-DUSE_FLOAT
CXXFLAGS+=-DUSE_EIGEN_TANH
CXXFLAGS+=-lm
CXXFLAGS+=-funroll-loops
CXXFLAGS+=-march=native
CXXFLAGS+=-m64
CXXFLAGS+=-DEIGEN_DONT_PARALLELIZE
CXXFLAGS+=-DEIGEN_NO_DEBUG
CXXFLAGS+=-DEIGEN_NO_STATIC_ASSERT
CXXFLAGS+=-I$(EIGEN_LOCATION)
CXXFLAGS+=-fopenmp

SRCS=$(shell ls *.cpp)
OBJS=$(SRCS:.cpp=.o)

PROGRAM=start

all : $(BUILD_DIR) $(patsubst %,$(BUILD_DIR)/%,$(PROGRAM))

$(BUILD_DIR)/%.o : %.cpp
	$(CXX) -c $(CXXFLAGS) -o $@ $<

$(BUILD_DIR)/$(PROGRAM) : $(patsubst %,$(BUILD_DIR)/%,$(OBJS))
	$(CXX) $(CXXFLAGS) $(CXXLIBS) -o $@ $^
	mv $(BUILD_DIR)/$(PROGRAM) ./
	rm -f ?*~

clean:
	rm -f $(BUILD_DIR)/* $(PROGRAM) ?*~
	echo "dummy" > objs/dummy

Notice if not using Makefile to build .cpp, then you can use -I flag to include paths (for current directory use: -I .), sometimes use -D flag for extra compiler option.

The Original training script (official)

Training scripts for custom network for torch C++ :
In main.cpp:

#include <torch/torch.h>
#include <iostream>
#include <iomanip>
#include "neural_net.h"

int main() {
   
    std::cout << "FeedForward Neural Network\n\n";

    // Device
    auto cuda_available = torch::cuda::is_available();
    torch::Device device(cuda_available ? torch::kCUDA : torch::kCPU);
    std::cout << (cuda_available ? "CUDA available. Training on GPU." : "Training on CPU.") << '\n';

    // Hyper parameters
    const int64_t input_size = 784;
    const int64_t hidden_size = 500;
    const int64_t num_classes = 10;
    const int64_t batch_size = 100;
    const size_t num_epochs = 5;
    const double learning_rate = 0.001;

    const std::string MNIST_data_path = "./data";

    // MNIST Dataset
    auto train_dataset = torch::data::datasets