pywarm
A cleaner way to build neural networks for PyTorch.
Examples | Tutorial | API reference
PyWarm is a lightweight, high-level neural network construction API for PyTorch. It enables defining all parts of NNs in the functional way.
With PyWarm, you can put all network data flow logic in the forward()
method of your model, without the need to define children modules in the __init__()
method and then call it again in the forward()
. This result in a much more readable model definition in fewer lines of code.
PyWarm only aims to simplify the network definition, and does not attempt to cover model training, validation or data handling.
For example, a convnet for MNIST: (If needed, click the tabs to switch between Warm and Torch versions)
# powered by PyWarmimport torch.nn as nnimport torch.nn.functional as Fimport warmimport warm.functional as Wclass ConvNet(nn.Module): def __init__(self): super().__init__() warm.up(self, [2, 1, 28, 28]) def forward(self, x): x = W.conv(x, 20, 5, activation='relu') x = F.max_pool2d(x, 2) x = W.conv(x, 50, 5, activation='relu') x = F.max_pool2d(x, 2) x = x.view(-1, 800) x = W.linear(x, 500, activation='relu') x = W.linear(x, 10) return F.log_softmax(x, dim=1)
# vanilla PyTorch version, taken from# pytorch tutorials/beginner_source/blitz/neural_networks_tutorial.py import torch.nn as nnimport torch.nn.functional as Fclass ConvNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(1, 20, 5, 1) self.conv2 = nn.Conv2d(20, 50, 5, 1) self.fc1 = nn.Linear(4*4*50, 500) self.fc2 = nn.Linear(500, 10) def forward(self, x): x = F.relu(self.conv1(x)) x = F.max_pool2d(x, 2, 2) x = F.relu(self.conv2(x)) x = F.max_pool2d(x, 2, 2) x = x.view(-1, 4*4*50) x = F.relu(self.fc1(x)) x = self.fc2(x) return F.log_softmax(x, dim=1)
A couple of things you may have noticed:
First of all, in the PyWarm version, the entire network definition and data flow logic resides in the forward()
method. You don't have to look up and down repeatedly to understand what self.conv1
, self.fc1
etc. is doing.
You do not need to track and specify in_channels
(or in_features
, etc.) for network layers. PyWarm can infer the information for you. e.g.
# Warmx = W.conv(x, 20, 5, activation='relu') x = W.conv(x, 50, 5, activation='relu')# Torchself.conv1 = nn.Conv2d(1, 20, 5, 1)self.conv2 = nn.Conv2d(20, 50, 5, 1)
One unified W.conv
for all 1D, 2D, and 3D cases. Fewer things to keep track of!
activation='relu'
. All warm.functional
APIs accept an optional activation
keyword, which is basically equivalent to F.relu(W.conv(...))
. The keyword activation
can also take in a callable, for example activation=torch.nn.ReLU(inplace=True)
or activation=swish
.
For deeper neural networks, see additional examples.
pip3 install pywarm
If you already have experinces with PyTorch, using PyWarm is very straightforward:
First, import PyWarm in you model file:
import warmimport warm.functional as W
Second, remove child module definitions in the model's __init__()
method. In stead, use W.conv
, W.linear
... etc. in the model's forward()
method, just like how you would use torch nn functional F.max_pool2d
, F.relu
... etc.
For example, instead of writing:
# Torchclass MyModule(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size) # other child module definitions def forward(self, x): x = self.conv1(x) # more forward steps
You can now write in the warm way:
# Warmclass MyWarmModule(nn.Module): def __init__(self): super().__init__() warm.up(self, input_shape_or_data) def forward(self, x): x = W.conv(x, out_channels, kernel_size) # no in_channels needed # more forward steps
Finally, don't forget to warmify the model by adding
warm.up(self, input_shape_or_data)
at the end of the model's __init__()
method. You need to supply input_shape_or_data
, which is either a tensor of input data, or just its shape, e.g. [2, 1, 28, 28]
for MNIST inputs.
The model is now ready to use, just like any other PyTorch models.
Check out the tutorial and examples if you want to learn more!
Clone the repository first, then
cd pywarm pytest -v
Documentations are generated using the excellent Portray package.
上一篇:rlpyt
下一篇:learn2learn
还没有评论,说两句吧!
热门资源
Keras-ResNeXt
Keras ResNeXt Implementation of ResNeXt models...
seetafaceJNI
项目介绍 基于中科院seetaface2进行封装的JAVA...
spark-corenlp
This package wraps Stanford CoreNLP annotators ...
capsnet-with-caps...
CapsNet with capsule-wise convolution Project ...
inferno-boilerplate
This is a very basic boilerplate example for pe...
智能在线
400-630-6780
聆听.建议反馈
E-mail: support@tusaishared.com