Restructuring C++ project:

Summary: * Reduce unnecessary header inclusions in models and io. * Move autocast to separate folder and hide autograd implementation in an anonymous namespace. * Moving files in subfolders. Reviewed By: fmassa Differential Revision: D25461523 fbshipit-source-id: 756eeb6848aacaa474de4825ed4c1045d17e2cea
pytorch · Dec 10, 2020 · eed0807 · eed0807
1 parent a0b13ca
commit eed0807
Show file tree

Hide file tree

Showing 104 changed files with 318 additions and 302 deletions.
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -28,19 +28,6 @@ function(CUDA_CONVERT_FLAGS EXISTING_TARGET)
     endif()
 endfunction()
 
-file(GLOB HEADERS torchvision/csrc/*.h)
-# Image extension
-file(GLOB IMAGE_HEADERS torchvision/csrc/cpu/image/*.h)
-file(GLOB IMAGE_SOURCES torchvision/csrc/cpu/image/*.cpp)
-file(GLOB OPERATOR_HEADERS torchvision/csrc/cpu/*.h)
-file(GLOB OPERATOR_SOURCES ${OPERATOR_HEADERS} torchvision/csrc/cpu/*.cpp ${IMAGE_HEADERS} ${IMAGE_SOURCES} ${HEADERS} torchvision/csrc/*.cpp)
-if(WITH_CUDA)
-  file(GLOB OPERATOR_HEADERS ${OPERATOR_HEADERS} torchvision/csrc/cuda/*.h)
-  file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} ${OPERATOR_HEADERS} torchvision/csrc/cuda/*.cu)
-endif()
-file(GLOB MODELS_HEADERS torchvision/csrc/models/*.h)
-file(GLOB MODELS_SOURCES torchvision/csrc/models/*.h torchvision/csrc/models/*.cpp)
-
 if(MSVC)
   set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /wd4819")
   if(WITH_CUDA)
@@ -64,7 +51,17 @@ endif()
 include(GNUInstallDirs)
 include(CMakePackageConfigHelpers)
 
-add_library(${PROJECT_NAME} SHARED ${MODELS_SOURCES} ${OPERATOR_SOURCES} ${IMAGE_SOURCES})
+set(TVCPP torchvision/csrc)
+list(APPEND ALLOW_LISTED ${TVCPP} ${TVCPP}/io/image ${TVCPP}/io/image/cpu ${TVCPP}/models ${TVCPP}/ops ${TVCPP}/ops/cpu)
+if(WITH_CUDA)
+    list(APPEND ALLOW_LISTED ${TVCPP}/ops/cuda ${TVCPP}/ops/autocast)
+endif()
+
+FOREACH(DIR ${ALLOW_LISTED})
+    file(GLOB ALL_SOURCES ${ALL_SOURCES} ${DIR}/*.*)
+ENDFOREACH()
+
+add_library(${PROJECT_NAME} SHARED ${ALL_SOURCES})
 target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES} ${PNG_LIBRARY} ${JPEG_LIBRARIES} Python3::Python)
 set_target_properties(${PROJECT_NAME} PROPERTIES
   EXPORT_NAME TorchVision
@@ -95,13 +92,8 @@ install(EXPORT TorchVisionTargets
   NAMESPACE TorchVision::
   DESTINATION ${TORCHVISION_CMAKECONFIG_INSTALL_DIR})
 
-install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME})
-install(FILES
-  ${OPERATOR_HEADERS}
-  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cpu)
-if(WITH_CUDA)
-  install(FILES
-    ${OPERATOR_HEADERS}
-    DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cuda)
-endif()
-install(FILES ${MODELS_HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/models)
+FOREACH(INPUT_DIR ${ALLOW_LISTED})
+    string(REPLACE "${TVCPP}" "${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}" OUTPUT_DIR ${INPUT_DIR})
+    file(GLOB INPUT_FILES ${INPUT_DIR}/*.*)
+    install(FILES ${INPUT_FILES} DESTINATION ${OUTPUT_DIR})
+ENDFOREACH()
diff --git a/examples/cpp/hello_world/main.cpp b/examples/cpp/hello_world/main.cpp
@@ -1,4 +1,6 @@
 #include <iostream>
+#include <torch/torch.h>
+#include <torchvision/vision.h>
 #include <torchvision/models/resnet.h>
 
 int main()

diff --git a/setup.py b/setup.py
@@ -134,8 +134,9 @@ def get_extensions():
     this_dir = os.path.dirname(os.path.abspath(__file__))
     extensions_dir = os.path.join(this_dir, 'torchvision', 'csrc')
 
-    main_file = glob.glob(os.path.join(extensions_dir, '*.cpp'))
-    source_cpu = glob.glob(os.path.join(extensions_dir, 'cpu', '*.cpp'))
+    main_file = glob.glob(os.path.join(extensions_dir, '*.cpp')) + glob.glob(os.path.join(extensions_dir, 'ops',
+                                                                                          '*.cpp'))
+    source_cpu = glob.glob(os.path.join(extensions_dir, 'ops', 'cpu', '*.cpp'))
 
     is_rocm_pytorch = False
     if torch.__version__ >= '1.5':
@@ -146,17 +147,19 @@ def get_extensions():
         hipify_python.hipify(
             project_directory=this_dir,
             output_directory=this_dir,
-            includes="torchvision/csrc/cuda/*",
+            includes="torchvision/csrc/ops/cuda/*",
             show_detailed=True,
             is_pytorch_extension=True,
         )
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'hip', '*.hip'))
+        source_cuda = glob.glob(os.path.join(extensions_dir, 'ops', 'hip', '*.hip'))
         # Copy over additional files
-        for file in glob.glob(r"torchvision/csrc/cuda/*.h"):
-            shutil.copy(file, "torchvision/csrc/hip")
+        for file in glob.glob(r"torchvision/csrc/ops/cuda/*.h"):
+            shutil.copy(file, "torchvision/csrc/ops/hip")
 
     else:
-        source_cuda = glob.glob(os.path.join(extensions_dir, 'cuda', '*.cu'))
+        source_cuda = glob.glob(os.path.join(extensions_dir, 'ops', 'cuda', '*.cu'))
+
+    source_cuda += glob.glob(os.path.join(extensions_dir, 'ops', 'autocast', '*.cpp'))
 
     sources = main_file + source_cpu
     extension = CppExtension
@@ -309,8 +312,8 @@ def get_extensions():
             image_library += [jpeg_lib]
             image_include += [jpeg_include]
 
-    image_path = os.path.join(extensions_dir, 'cpu', 'image')
-    image_src = glob.glob(os.path.join(image_path, '*.cpp'))
+    image_path = os.path.join(extensions_dir, 'io', 'image')
+    image_src = glob.glob(os.path.join(image_path, '*.cpp')) + glob.glob(os.path.join(image_path, 'cpu', '*.cpp'))
 
     if png_found or jpeg_found:
         ext_modules.append(extension(
@@ -377,13 +380,13 @@ def get_extensions():
         print("ffmpeg library_dir: {}".format(ffmpeg_library_dir))
 
         # TorchVision base decoder + video reader
-        video_reader_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'video_reader')
+        video_reader_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'video_reader')
         video_reader_src = glob.glob(os.path.join(video_reader_src_dir, "*.cpp"))
-        base_decoder_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'decoder')
+        base_decoder_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'decoder')
         base_decoder_src = glob.glob(
             os.path.join(base_decoder_src_dir, "*.cpp"))
         # Torchvision video API
-        videoapi_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'cpu', 'video')
+        videoapi_src_dir = os.path.join(this_dir, 'torchvision', 'csrc', 'io', 'video')
         videoapi_src = glob.glob(os.path.join(videoapi_src_dir, "*.cpp"))
         # exclude tests
         base_decoder_src = [x for x in base_decoder_src if '_test.cpp' not in x]

diff --git a/test/tracing/frcnn/test_frcnn_tracing.cpp b/test/tracing/frcnn/test_frcnn_tracing.cpp
@@ -1,8 +1,7 @@
-#include <ATen/ATen.h>
 #include <torch/script.h>
 #include <torch/torch.h>
-#include <torchvision/nms.h>
-#include <torchvision/roi_align.h>
+#include <torchvision/vision.h>
+#include <torchvision/ops/nms.h>
 
 #ifdef _WIN32
 // Windows only

diff --git a/torchvision/csrc/cpu/image/image.h b/torchvision/csrc/cpu/image/image.h
diff --git a/...vision/csrc/cpu/decoder/audio_sampler.cpp → ...hvision/csrc/io/decoder/audio_sampler.cpp b/...vision/csrc/cpu/decoder/audio_sampler.cpp → ...hvision/csrc/io/decoder/audio_sampler.cpp
diff --git a/torchvision/csrc/cpu/decoder/audio_sampler.h → torchvision/csrc/io/decoder/audio_sampler.h b/torchvision/csrc/cpu/decoder/audio_sampler.h → torchvision/csrc/io/decoder/audio_sampler.h
diff --git a/...hvision/csrc/cpu/decoder/audio_stream.cpp → torchvision/csrc/io/decoder/audio_stream.cpp b/...hvision/csrc/cpu/decoder/audio_stream.cpp → torchvision/csrc/io/decoder/audio_stream.cpp
diff --git a/torchvision/csrc/cpu/decoder/audio_stream.h → torchvision/csrc/io/decoder/audio_stream.h b/torchvision/csrc/cpu/decoder/audio_stream.h → torchvision/csrc/io/decoder/audio_stream.h
diff --git a/torchvision/csrc/cpu/decoder/cc_stream.cpp → torchvision/csrc/io/decoder/cc_stream.cpp b/torchvision/csrc/cpu/decoder/cc_stream.cpp → torchvision/csrc/io/decoder/cc_stream.cpp
diff --git a/torchvision/csrc/cpu/decoder/cc_stream.h → torchvision/csrc/io/decoder/cc_stream.h b/torchvision/csrc/cpu/decoder/cc_stream.h → torchvision/csrc/io/decoder/cc_stream.h
diff --git a/torchvision/csrc/cpu/decoder/decoder.cpp → torchvision/csrc/io/decoder/decoder.cpp b/torchvision/csrc/cpu/decoder/decoder.cpp → torchvision/csrc/io/decoder/decoder.cpp
diff --git a/torchvision/csrc/cpu/decoder/decoder.h → torchvision/csrc/io/decoder/decoder.h b/torchvision/csrc/cpu/decoder/decoder.h → torchvision/csrc/io/decoder/decoder.h
diff --git a/torchvision/csrc/cpu/decoder/defs.h → torchvision/csrc/io/decoder/defs.h b/torchvision/csrc/cpu/decoder/defs.h → torchvision/csrc/io/decoder/defs.h
diff --git a/...vision/csrc/cpu/decoder/memory_buffer.cpp → ...hvision/csrc/io/decoder/memory_buffer.cpp b/...vision/csrc/cpu/decoder/memory_buffer.cpp → ...hvision/csrc/io/decoder/memory_buffer.cpp
diff --git a/torchvision/csrc/cpu/decoder/memory_buffer.h → torchvision/csrc/io/decoder/memory_buffer.h b/torchvision/csrc/cpu/decoder/memory_buffer.h → torchvision/csrc/io/decoder/memory_buffer.h
diff --git a/...sion/csrc/cpu/decoder/seekable_buffer.cpp → ...ision/csrc/io/decoder/seekable_buffer.cpp b/...sion/csrc/cpu/decoder/seekable_buffer.cpp → ...ision/csrc/io/decoder/seekable_buffer.cpp
diff --git a/...vision/csrc/cpu/decoder/seekable_buffer.h → ...hvision/csrc/io/decoder/seekable_buffer.h b/...vision/csrc/cpu/decoder/seekable_buffer.h → ...hvision/csrc/io/decoder/seekable_buffer.h
diff --git a/torchvision/csrc/cpu/decoder/stream.cpp → torchvision/csrc/io/decoder/stream.cpp b/torchvision/csrc/cpu/decoder/stream.cpp → torchvision/csrc/io/decoder/stream.cpp
diff --git a/torchvision/csrc/cpu/decoder/stream.h → torchvision/csrc/io/decoder/stream.h b/torchvision/csrc/cpu/decoder/stream.h → torchvision/csrc/io/decoder/stream.h
diff --git a/...ion/csrc/cpu/decoder/subtitle_sampler.cpp → ...sion/csrc/io/decoder/subtitle_sampler.cpp b/...ion/csrc/cpu/decoder/subtitle_sampler.cpp → ...sion/csrc/io/decoder/subtitle_sampler.cpp
diff --git a/...ision/csrc/cpu/decoder/subtitle_sampler.h → ...vision/csrc/io/decoder/subtitle_sampler.h b/...ision/csrc/cpu/decoder/subtitle_sampler.h → ...vision/csrc/io/decoder/subtitle_sampler.h
diff --git a/...sion/csrc/cpu/decoder/subtitle_stream.cpp → ...ision/csrc/io/decoder/subtitle_stream.cpp b/...sion/csrc/cpu/decoder/subtitle_stream.cpp → ...ision/csrc/io/decoder/subtitle_stream.cpp
diff --git a/...vision/csrc/cpu/decoder/subtitle_stream.h → ...hvision/csrc/io/decoder/subtitle_stream.h b/...vision/csrc/cpu/decoder/subtitle_stream.h → ...hvision/csrc/io/decoder/subtitle_stream.h
diff --git a/...hvision/csrc/cpu/decoder/sync_decoder.cpp → torchvision/csrc/io/decoder/sync_decoder.cpp b/...hvision/csrc/cpu/decoder/sync_decoder.cpp → torchvision/csrc/io/decoder/sync_decoder.cpp
diff --git a/torchvision/csrc/cpu/decoder/sync_decoder.h → torchvision/csrc/io/decoder/sync_decoder.h b/torchvision/csrc/cpu/decoder/sync_decoder.h → torchvision/csrc/io/decoder/sync_decoder.h
diff --git a/...on/csrc/cpu/decoder/sync_decoder_test.cpp → ...ion/csrc/io/decoder/sync_decoder_test.cpp b/...on/csrc/cpu/decoder/sync_decoder_test.cpp → ...ion/csrc/io/decoder/sync_decoder_test.cpp
diff --git a/torchvision/csrc/cpu/decoder/time_keeper.cpp → torchvision/csrc/io/decoder/time_keeper.cpp b/torchvision/csrc/cpu/decoder/time_keeper.cpp → torchvision/csrc/io/decoder/time_keeper.cpp
diff --git a/torchvision/csrc/cpu/decoder/time_keeper.h → torchvision/csrc/io/decoder/time_keeper.h b/torchvision/csrc/cpu/decoder/time_keeper.h → torchvision/csrc/io/decoder/time_keeper.h
diff --git a/torchvision/csrc/cpu/decoder/util.cpp → torchvision/csrc/io/decoder/util.cpp b/torchvision/csrc/cpu/decoder/util.cpp → torchvision/csrc/io/decoder/util.cpp
diff --git a/torchvision/csrc/cpu/decoder/util.h → torchvision/csrc/io/decoder/util.h b/torchvision/csrc/cpu/decoder/util.h → torchvision/csrc/io/decoder/util.h
diff --git a/torchvision/csrc/cpu/decoder/util_test.cpp → torchvision/csrc/io/decoder/util_test.cpp b/torchvision/csrc/cpu/decoder/util_test.cpp → torchvision/csrc/io/decoder/util_test.cpp
diff --git a/...vision/csrc/cpu/decoder/video_sampler.cpp → ...hvision/csrc/io/decoder/video_sampler.cpp b/...vision/csrc/cpu/decoder/video_sampler.cpp → ...hvision/csrc/io/decoder/video_sampler.cpp
diff --git a/torchvision/csrc/cpu/decoder/video_sampler.h → torchvision/csrc/io/decoder/video_sampler.h b/torchvision/csrc/cpu/decoder/video_sampler.h → torchvision/csrc/io/decoder/video_sampler.h
diff --git a/...hvision/csrc/cpu/decoder/video_stream.cpp → torchvision/csrc/io/decoder/video_stream.cpp b/...hvision/csrc/cpu/decoder/video_stream.cpp → torchvision/csrc/io/decoder/video_stream.cpp
diff --git a/torchvision/csrc/cpu/decoder/video_stream.h → torchvision/csrc/io/decoder/video_stream.h b/torchvision/csrc/cpu/decoder/video_stream.h → torchvision/csrc/io/decoder/video_stream.h
diff --git a/...hvision/csrc/cpu/image/read_image_cpu.cpp → ...sion/csrc/io/image/cpu/read_image_cpu.cpp b/...hvision/csrc/cpu/image/read_image_cpu.cpp → ...sion/csrc/io/image/cpu/read_image_cpu.cpp
@@ -1,4 +1,5 @@
 #include "read_image_cpu.h"
+
 #include "readjpeg_cpu.h"
 #include "readpng_cpu.h"
 

diff --git a/torchvision/csrc/cpu/image/read_image_cpu.h → ...vision/csrc/io/image/cpu/read_image_cpu.h b/torchvision/csrc/cpu/image/read_image_cpu.h → ...vision/csrc/io/image/cpu/read_image_cpu.h
@@ -1,7 +1,7 @@
 #pragma once
 
-#include <torch/torch.h>
-#include "image_read_mode.h"
+#include <torch/types.h>
+#include "../image_read_mode.h"
 
 C10_EXPORT torch::Tensor decode_image(
     const torch::Tensor& data,

diff --git a/...on/csrc/cpu/image/read_write_file_cpu.cpp → ...csrc/io/image/cpu/read_write_file_cpu.cpp b/...on/csrc/cpu/image/read_write_file_cpu.cpp → ...csrc/io/image/cpu/read_write_file_cpu.cpp
diff --git a/...sion/csrc/cpu/image/read_write_file_cpu.h → ...n/csrc/io/image/cpu/read_write_file_cpu.h b/...sion/csrc/cpu/image/read_write_file_cpu.h → ...n/csrc/io/image/cpu/read_write_file_cpu.h
@@ -1,8 +1,7 @@
 #pragma once
 
-#include <errno.h>
 #include <sys/stat.h>
-#include <torch/torch.h>
+#include <torch/types.h>
 
 C10_EXPORT torch::Tensor read_file(const std::string& filename);
 

diff --git a/torchvision/csrc/cpu/image/readjpeg_cpu.cpp → ...vision/csrc/io/image/cpu/readjpeg_cpu.cpp b/torchvision/csrc/cpu/image/readjpeg_cpu.cpp → ...vision/csrc/io/image/cpu/readjpeg_cpu.cpp
@@ -1,16 +1,12 @@
 #include "readjpeg_cpu.h"
 
-#include <ATen/ATen.h>
-
 #if !JPEG_FOUND
 torch::Tensor decodeJPEG(const torch::Tensor& data, ImageReadMode mode) {
   TORCH_CHECK(
       false, "decodeJPEG: torchvision not compiled with libjpeg support");
 }
 #else
-#include <jpeglib.h>
-#include <setjmp.h>
-#include "jpegcommon.h"
+#include "../jpegcommon.h"
 
 struct torch_jpeg_mgr {
   struct jpeg_source_mgr pub;

diff --git a/torchvision/csrc/cpu/image/readjpeg_cpu.h → torchvision/csrc/io/image/cpu/readjpeg_cpu.h b/torchvision/csrc/cpu/image/readjpeg_cpu.h → torchvision/csrc/io/image/cpu/readjpeg_cpu.h
@@ -1,7 +1,7 @@
 #pragma once
 
-#include <torch/torch.h>
-#include "image_read_mode.h"
+#include <torch/types.h>
+#include "../image_read_mode.h"
 
 C10_EXPORT torch::Tensor decodeJPEG(
     const torch::Tensor& data,

diff --git a/torchvision/csrc/cpu/image/readpng_cpu.cpp → ...hvision/csrc/io/image/cpu/readpng_cpu.cpp b/torchvision/csrc/cpu/image/readpng_cpu.cpp → ...hvision/csrc/io/image/cpu/readpng_cpu.cpp
@@ -1,7 +1,5 @@
 #include "readpng_cpu.h"
 
-#include <ATen/ATen.h>
-
 #if !PNG_FOUND
 torch::Tensor decodePNG(const torch::Tensor& data, ImageReadMode mode) {
   TORCH_CHECK(false, "decodePNG: torchvision not compiled with libPNG support");

diff --git a/torchvision/csrc/cpu/image/readpng_cpu.h → torchvision/csrc/io/image/cpu/readpng_cpu.h b/torchvision/csrc/cpu/image/readpng_cpu.h → torchvision/csrc/io/image/cpu/readpng_cpu.h
@@ -1,7 +1,7 @@
 #pragma once
 
-#include <torch/torch.h>
-#include "image_read_mode.h"
+#include <torch/types.h>
+#include "../image_read_mode.h"
 
 C10_EXPORT torch::Tensor decodePNG(
     const torch::Tensor& data,

diff --git a/torchvision/csrc/cpu/image/writejpeg_cpu.cpp → ...ision/csrc/io/image/cpu/writejpeg_cpu.cpp b/torchvision/csrc/cpu/image/writejpeg_cpu.cpp → ...ision/csrc/io/image/cpu/writejpeg_cpu.cpp
@@ -1,8 +1,5 @@
 #include "writejpeg_cpu.h"
 
-#include <setjmp.h>
-#include <string>
-
 #if !JPEG_FOUND
 
 torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality) {
@@ -11,9 +8,7 @@ torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality) {
 }
 
 #else
-
-#include <jpeglib.h>
-#include "jpegcommon.h"
+#include "../jpegcommon.h"
 
 torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality) {
   // Define compression structures and error handling

diff --git a/torchvision/csrc/cpu/image/writejpeg_cpu.h → ...hvision/csrc/io/image/cpu/writejpeg_cpu.h b/torchvision/csrc/cpu/image/writejpeg_cpu.h → ...hvision/csrc/io/image/cpu/writejpeg_cpu.h
@@ -1,5 +1,5 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/types.h>
 
 C10_EXPORT torch::Tensor encodeJPEG(const torch::Tensor& data, int64_t quality);
diff --git a/torchvision/csrc/cpu/image/writepng_cpu.cpp → ...vision/csrc/io/image/cpu/writepng_cpu.cpp b/torchvision/csrc/cpu/image/writepng_cpu.cpp → ...vision/csrc/io/image/cpu/writepng_cpu.cpp
@@ -1,17 +1,14 @@
 #include "writejpeg_cpu.h"
 
-#include <setjmp.h>
-#include <string>
-
 #if !PNG_FOUND
 
 torch::Tensor encodePNG(const torch::Tensor& data, int64_t compression_level) {
   TORCH_CHECK(false, "encodePNG: torchvision not compiled with libpng support");
 }
 
 #else
-
 #include <png.h>
+#include <setjmp.h>
 
 struct torch_mem_encode {
   char* buffer;

diff --git a/torchvision/csrc/cpu/image/writepng_cpu.h → torchvision/csrc/io/image/cpu/writepng_cpu.h b/torchvision/csrc/cpu/image/writepng_cpu.h → torchvision/csrc/io/image/cpu/writepng_cpu.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/types.h>
 
 C10_EXPORT torch::Tensor encodePNG(
     const torch::Tensor& data,

diff --git a/torchvision/csrc/cpu/image/image.cpp → torchvision/csrc/io/image/image.cpp b/torchvision/csrc/cpu/image/image.cpp → torchvision/csrc/io/image/image.cpp
@@ -1,6 +1,5 @@
-
 #include "image.h"
-#include <ATen/ATen.h>
+
 #include <Python.h>
 
 // If we are in a Windows environment, we need to define

diff --git a/torchvision/csrc/io/image/image.h b/torchvision/csrc/io/image/image.h
@@ -0,0 +1,8 @@
+#pragma once
+
+#include "cpu/read_image_cpu.h"
+#include "cpu/read_write_file_cpu.h"
+#include "cpu/readjpeg_cpu.h"
+#include "cpu/readpng_cpu.h"
+#include "cpu/writejpeg_cpu.h"
+#include "cpu/writepng_cpu.h"
diff --git a/torchvision/csrc/cpu/image/image_read_mode.h → torchvision/csrc/io/image/image_read_mode.h b/torchvision/csrc/cpu/image/image_read_mode.h → torchvision/csrc/io/image/image_read_mode.h
diff --git a/torchvision/csrc/cpu/image/jpegcommon.cpp → torchvision/csrc/io/image/jpegcommon.cpp b/torchvision/csrc/cpu/image/jpegcommon.cpp → torchvision/csrc/io/image/jpegcommon.cpp
@@ -1,7 +1,6 @@
+#if JPEG_FOUND
 #include "jpegcommon.h"
-#include <string>
 
-#if JPEG_FOUND
 void torch_jpeg_error_exit(j_common_ptr cinfo) {
   /* cinfo->err really points to a torch_jpeg_error_mgr struct, so coerce
    * pointer */

diff --git a/torchvision/csrc/cpu/image/jpegcommon.h → torchvision/csrc/io/image/jpegcommon.h b/torchvision/csrc/cpu/image/jpegcommon.h → torchvision/csrc/io/image/jpegcommon.h
@@ -8,7 +8,6 @@
 #if JPEG_FOUND
 #include <jpeglib.h>
 #include <setjmp.h>
-#include <string>
 
 static const JOCTET EOI_BUFFER[1] = {JPEG_EOI};
 struct torch_jpeg_error_mgr {

diff --git a/torchvision/csrc/cpu/video/Video.cpp → torchvision/csrc/io/video/Video.cpp b/torchvision/csrc/cpu/video/Video.cpp → torchvision/csrc/io/video/Video.cpp
diff --git a/torchvision/csrc/cpu/video/Video.h → torchvision/csrc/io/video/Video.h b/torchvision/csrc/cpu/video/Video.h → torchvision/csrc/io/video/Video.h
diff --git a/torchvision/csrc/cpu/video/register.cpp → torchvision/csrc/io/video/register.cpp b/torchvision/csrc/cpu/video/register.cpp → torchvision/csrc/io/video/register.cpp
diff --git a/...ion/csrc/cpu/video_reader/VideoReader.cpp → ...sion/csrc/io/video_reader/VideoReader.cpp b/...ion/csrc/cpu/video_reader/VideoReader.cpp → ...sion/csrc/io/video_reader/VideoReader.cpp
diff --git a/...ision/csrc/cpu/video_reader/VideoReader.h → ...vision/csrc/io/video_reader/VideoReader.h b/...ision/csrc/cpu/video_reader/VideoReader.h → ...vision/csrc/io/video_reader/VideoReader.h
diff --git a/torchvision/csrc/models/alexnet.h b/torchvision/csrc/models/alexnet.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {

diff --git a/torchvision/csrc/models/densenet.h b/torchvision/csrc/models/densenet.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {

diff --git a/torchvision/csrc/models/googlenet.cpp b/torchvision/csrc/models/googlenet.cpp
@@ -1,7 +1,5 @@
 #include "googlenet.h"
 
-#include "modelsimpl.h"
-
 namespace vision {
 namespace models {
 

diff --git a/torchvision/csrc/models/googlenet.h b/torchvision/csrc/models/googlenet.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {

diff --git a/torchvision/csrc/models/inception.h b/torchvision/csrc/models/inception.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {

diff --git a/torchvision/csrc/models/mnasnet.h b/torchvision/csrc/models/mnasnet.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {

diff --git a/torchvision/csrc/models/mobilenet.h b/torchvision/csrc/models/mobilenet.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {

diff --git a/torchvision/csrc/models/modelsimpl.h b/torchvision/csrc/models/modelsimpl.h
@@ -1,10 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
-
-#ifndef TORCH_CHECK
-#define TORCH_CHECK AT_CHECK
-#endif
+#include <torch/nn.h>
 
 namespace vision {
 namespace models {

diff --git a/torchvision/csrc/models/resnet.cpp b/torchvision/csrc/models/resnet.cpp
@@ -1,7 +1,5 @@
 #include "resnet.h"
 
-#include "modelsimpl.h"
-
 namespace vision {
 namespace models {
 namespace _resnetimpl {

diff --git a/torchvision/csrc/models/resnet.h b/torchvision/csrc/models/resnet.h
@@ -1,6 +1,6 @@
 #pragma once
 
-#include <torch/torch.h>
+#include <torch/nn.h>
 #include "../macros.h"
 
 namespace vision {