| # Copyright (c) Chris Choy (chrischoy@ai.stanford.edu). | |
| # | |
| # Permission is hereby granted, free of charge, to any person obtaining a copy of | |
| # this software and associated documentation files (the "Software"), to deal in | |
| # the Software without restriction, including without limitation the rights to | |
| # use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies | |
| # of the Software, and to permit persons to whom the Software is furnished to do | |
| # so, subject to the following conditions: | |
| # | |
| # The above copyright notice and this permission notice shall be included in all | |
| # copies or substantial portions of the Software. | |
| # | |
| # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR | |
| # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, | |
| # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE | |
| # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER | |
| # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, | |
| # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE | |
| # SOFTWARE. | |
| # | |
| # Please cite "4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural | |
| # Networks", CVPR'19 (https://arxiv.org/abs/1904.08755) if you use any part | |
| # of the code. | |
| import torch | |
| import torch.nn as nn | |
| from torch.optim import SGD | |
| import MinkowskiEngine as ME | |
| from tests.python.common import data_loader | |
| class ExampleNetwork(ME.MinkowskiNetwork): | |
| def __init__(self, in_feat, out_feat, D): | |
| super(ExampleNetwork, self).__init__(D) | |
| self.net = nn.Sequential( | |
| ME.MinkowskiConvolution( | |
| in_channels=in_feat, | |
| out_channels=64, | |
| kernel_size=3, | |
| stride=2, | |
| dilation=1, | |
| bias=False, | |
| dimension=D), ME.MinkowskiBatchNorm(64), ME.MinkowskiReLU(), | |
| ME.MinkowskiConvolution( | |
| in_channels=64, | |
| out_channels=128, | |
| kernel_size=3, | |
| stride=2, | |
| dimension=D), ME.MinkowskiBatchNorm(128), ME.MinkowskiReLU(), | |
| ME.MinkowskiGlobalPooling(), | |
| ME.MinkowskiLinear(128, out_feat)) | |
| def forward(self, x): | |
| return self.net(x) | |
| if __name__ == '__main__': | |
| # loss and network | |
| criterion = nn.CrossEntropyLoss() | |
| net = ExampleNetwork(in_feat=3, out_feat=5, D=2) | |
| print(net) | |
| # a data loader must return a tuple of coords, features, and labels. | |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
| net = net.to(device) | |
| optimizer = SGD(net.parameters(), lr=1e-1) | |
| for i in range(10): | |
| optimizer.zero_grad() | |
| # Get new data | |
| coords, feat, label = data_loader() | |
| input = ME.SparseTensor(feat, coords, device=device) | |
| label = label.to(device) | |
| # Forward | |
| output = net(input) | |
| # Loss | |
| loss = criterion(output.F, label) | |
| print('Iteration: ', i, ', Loss: ', loss.item()) | |
| # Gradient | |
| loss.backward() | |
| optimizer.step() | |
| # Saving and loading a network | |
| torch.save(net.state_dict(), 'test.pth') | |
| net.load_state_dict(torch.load('test.pth')) | |