defModule macro to define neural nets without inline C++

[mratsim]

Currently we need to inline C++ code due to Nim not being able to generate C++ code with default value or wrap types without default constructors nim-lang/Nim#4687

This leads to the following code

flambeau/proof_of_concepts/poc09_end_to_end.nim

Lines 12 to 26 in fe82c68

	emitTypes:
	"""
	struct Net: public torch::nn::Module {
	torch::nn::Linear fc1{nullptr};
	torch::nn::Linear fc2{nullptr};
	torch::nn::Linear fc3{nullptr};
	};
	"""
	type Net
	{.pure, importcpp.}
	= object of Module
	fc1: Linear
	fc2: Linear
	fc3: Linear

with emittypes a workaround nim-lang/Nim#16664

flambeau/flambeau/cpp/emitters.nim

Lines 3 to 13 in fe82c68

	template emitTypes*(typesection: static string): untyped =
	## Emit a C/C++ typesection
	{.emit: "/*TYPESECTION*/\n" & typesection.}
	template emitGlobals*(globals: static string): untyped =
	## Emit a C/C++ global variable declaration
	{.emit: "/*VARSECTION*/\n" & globals.}
	template emitIncludes*(includes: static string): untyped =
	## Emit a C/C++ global variable declaration
	{.emit: "/*INCLUDESECTION*/\n" & includes.}

Instead it would be more convenient to have a macro defModule (name subject to bikeshedding) with the following syntax:

defModule:
  type Net = object of Module
    fc1: Linear = nil
    fc2: Linear = nil
    fc3: Linear = nil

and generates the proper C++ code (nil mapped to {nullptr}) and rewrite the typesection with the proper {.pure, importcpp.} pragma.

This is similar to nim-lang/RFCs#126, nim-lang/RFCs#252

Implementation details:

• To interpolate with the proper symbol, use the syntax {.emit:["type something {", NimTypeInterpolated, " field{nullptr};"].}
• With interpolation we should properly support both C++ types without extracting from their importcpp pragma and Nim types without fiddling with gensym.

[HugoGranstrom]

Thus far I've got an untyped macro that can generate the emitted C++ code and add the correct pragmas. Would it make sense to also generate a init proc as well like the one here:

flambeau/examples/pytorch/cpp/mnist/mnist.nim

Lines 24 to 47 in cba9cc5

	emitTypes:
	"""
	struct Net: public torch::nn::Module {
	torch::nn::Conv2d conv1{nullptr};
	torch::nn::Conv2d conv2{nullptr};
	torch::nn::Dropout2d conv2_drop{nullptr};
	torch::nn::Linear fc1{nullptr};
	torch::nn::Linear fc2{nullptr};
	};
	"""
	type Net {.pure, importcpp.} = object of Module
	conv1: Conv2d
	conv2: Conv2d
	conv2_drop: Dropout2d
	fc1: Linear
	fc2: Linear
	proc init(net: var Net) =
	net.conv1 = net.register_module("conv1", Conv2d.init(1, 10, 5))
	net.conv2 = net.register_module("conv2", Conv2d.init(10, 20, 5))
	net.conv2_drop = net.register_module("conv2_drop", Dropout2d.init())
	net.fc1 = net.register_module("fc1", Linear.init(320, 50))
	net.fc2 = net.register_module("fc2", Linear.init(50, 10))

We could use your syntax from above (if we assume all types are initialized with nullptr) and write it like this:

defModule:
  type Net = object of Module 
   conv1: Conv2d = (1, 10, 5)
   conv2: Conv2d = (10, 20, 5)
   conv2_drop: Dropout2d # no value means passing no parameters
   fc1: Linear = (320, 50)
   fc2: Linear = (50, 10)

Then the values would be passed to the type's init proc. We could make it optional to generate it as well so people could choose. But this does of course build upon the fact that we are using an untyped macro as the syntax isn't legal Nim code.

I'm not too used to Torch yet so I'm not really sure whether this is an approach that would work in general for a typical Net, so any feedback is appreciated.

[mratsim]

I think something like this is better:

defModule:
  type Net = object of Module 
    conv1 = Conv2d(1, 10, 5)
    conv2 = Conv2d(10, 20, 5)
    conv2_drop: Dropout2d # no value means passing no parameters
    fc1 = Linear(320, 50)
    fc2 = Linear(50, 10)

It avoids floating untyped parenthesis and mimics object construction.

The tricky parts are:

• shared weights which can have a dedicated syntax
• nested nets, for example initializing a Net2 from Net

[HugoGranstrom]

That does indeed look more intuitive. 👍 And with the power of pattern matching not that much harder either 😎

• Yeah, I'll keep in mind! I'll try my best to make the code extendable for future special syntaxes.
• That one sounds a bit trickier 🤔 Are you thinking about transfer learning and using the weights from a pretrained network or something like that? Then I could see the problems if one wants to modify it by replacing the last layers or something.

And another thing that I think got lost yesterday, should we support Nim types in a Net? I don't really see how we could include them in the emitted C++ types so I'd say: no 😛 Plus we have no way of telling if it is a Nim type or C++ type in a untyped network.

[HugoGranstrom]

Think I realized what you meant by nested nets now, something along these lines, right?

defModule:
  type 
	Net = object of Module 
		conv1 = Conv2d(1, 10, 5)
		conv2 = Conv2d(10, 20, 5)
	    conv2_drop: Dropout2d # no value means passing no parameters
	    fc1 = Linear(320, 50)
	    fc2 = Linear(50, 10)
	Net2 = object of Module
		net1: Net1
		fc = Linear(100, 100)

If we go for an untyped macro we'd need a separate syntax for nested nets. Perhaps we could do it like this, that custom Nets are specified in "Nim-style" (net1: Net1) and built-in layers are defined using our new "Construction-style" (fc = Linear(100, 100)):

Net2 = object of Module
  net1: Net1
  fc = Linear(100, 100)

This would be easy to parse at least. But the again, it depends on what types we should allow in a Net.

[Vindaar]

Why not just do:

defModule:
  type 
    Net = object of Module 
      conv1 = Conv2d(1, 10, 5)
      conv2 = Conv2d(10, 20, 5)
      conv2_drop: Dropout2d # no value means passing no parameters
      fc1 = Linear(320, 50)
      fc2 = Linear(50, 10)
    Net2 = object of Module
      net1 = Net()
      fc = Linear(100, 100)

in those cases?

[HugoGranstrom]

@Vindaar Because we need to be able to differentiate between them when generating the emitted C++. The code for Net2 should be:

struct Net2: public torch::nn::Module { 
     Net net1{nullptr}; // here's the difference, no namespace
     torch::nn::Linear fc{nullptr}; 
   }; 

But the generated code will be this if we can't see a difference between custom Nets (that we generate the C++ code for) and the built-in ones:

struct Net2: public torch::nn::Module { 
     torch::nn::Net net1{nullptr}; // here's the difference, wrong namespace
     torch::nn::Linear fc{nullptr}; 
   }; 

Built-in layers like conv2_drop: Dropout2d could also miss initial parameters and could (and perhaps should) be specified as conv2_drop = Dropout2d(). Perhaps that clashes with what you suggested?

[Vindaar]

Ah, but that can be easily solved by keeping track of all nets that were user generated (e.g. Net and Net2 in this example) and only prefixing the torch:nn prefix for those, which are not user defined, no?

[HugoGranstrom]

Yes, that should work but then we're back at how to do that safely between macros. As people could use multiple defModule calls and use types defined in one in another. And thus we would need a shared CT table/set. (macrocache doesn't have support for checking if a key is in a table or not btw)

edit: This approach also makes the C++ generation more complex as it must first parse all the code, and then create all of the C++ structs. The current approach just creates the C++ code while parsing which is more simple. With different syntaxes that could still be the case, but that's just me being lazy 🤣

[Vindaar]

Oh ok, now I'm starting to see what you're talking about.

A shared CT table is something that should be supported anyways, imo. So it's only a matter of time until that is / should be implemented. Maybe one of us can help out improving the macrocache module for the purpose?

With the fully untyped, no state approach, another option might be something like wrapping the name in a call, e.g.:

defModule:
  type 
    Net = object of Module 
      conv1 = Conv2d(1, 10, 5)
      conv2 = Conv2d(10, 20, 5)
      conv2_drop: Dropout2d # no value means passing no parameters
      fc1 = Linear(320, 50)
      fc2 = Linear(50, 10)
    Net2 = object of Module
      net1 = custom(Net())
      fc = Linear(100, 100)

e.g. custom to indicate this is not a builtin torch type. Bikeshed on the name. custom, local, ...
Leaves the syntax alive, is easy to extract, can be documented well.

[HugoGranstrom]

Yeah, it definitely should be supported, and we probably should help improve it. But it feels like it's quite convoluted to handle everything in an IC-safe way without digging our fingers into how the internals of IC works 🤔

That approach is more feasible and future-proof in my opinion 😄 And you have a good point in that it is easier to document and understand than different syntaxes. And the hardest part is of course finding a good name for it 🤣 but custom is a good name that may very well stay. And it looks even more like a syntactic thing with net1 = custom Net() (that reminds me I'll have to handle two cases:P)

[HugoGranstrom]

Mnist now works without having to manually emit the C++ type and the init proc is generated automatically as well. I'm not really sure how to handle the nested Nets though, and specifically how to emit the C++ code for them. I thought I could just do:

struct Net: public torch::nn::Module {
	torch::nn::Conv2d conv1{nullptr};
	torch::nn::Conv2d conv2{nullptr};
	torch::nn::Dropout2d conv2_drop{nullptr};
	torch::nn::Linear fc1{nullptr};
	torch::nn::Linear fc2{nullptr};
};
struct Net2: public torch::nn::Module {
	torch::nn::Linear fc1{nullptr};
	Net net1{nullptr}; // error here
};

But then I get this error:

error C2664: 'Net::Net(const Net &)': cannot convert argument 1 from 'nullptr' to 'const Net &'
note: Reason: cannot convert from 'nullptr' to 'const Net'
note: No constructor could take the source type, or constructor overload resolution was ambiguous

It feels like Net isn't a shared_ptr really as it should be? Any ideas on how to approach it? Must we define it as a pointer ourselves or does the inheritance handle that for us?

[Vindaar]

Great work!

About the issue, maybe this doesn't make any sense, but:
Maybe the issue is simply that Net does not define a constructor that takes a single pointer? So that the solution would be to add a simple constructor line for each custom module we create? I tried checking the source of what kind of constructors those already implemented classes provide, but when I saw the TORCH_MODULE macro here:
https://pytorch.org/cppdocs/api/program_listing_file_torch_csrc_api_include_torch_nn_modules_linear.h.html
I kinda gave up.

To be honest, I don't even know exactly what providing a value within a struct definition does. I suppose a default value?

[HugoGranstrom]

Thanks 😄

That sound very reasonable, the built-in layers are shared_ptrs so making nullptr default for them makes sense but not for us. So the question is what to put in there instead. I had a look here https://pytorch.org/tutorials/advanced/cpp_frontend.html and they seem to define a struct first and then run the TORCH_MODULE to create a pointer-y type of it somehow. But I'm not sure if we want that? We would have to change all net.param to net->param, right? Plus they have the initialization in the struct body itself which sucks for us. I'm going to think about this in the sauna for a moment now 🤔

[HugoGranstrom]

Something as easy as removing the default value altogether seems to have stopped the compiler from complaining at least: Net net1;. Sometimes it's so easy 🤣 Don't know if it works though must test a nested net if it runs, but that's for tomorrow 🤞