Skip to content

Commit

Permalink
fix(): first commit
Browse files Browse the repository at this point in the history
Signed-off-by: Abhiram Iyer <abhi.iyer.ai@gmail.com>
Signed-off-by: Abhiram Iyer <abhirami@nvidia.com>
  • Loading branch information
abhi-iyer committed Jun 8, 2020
1 parent 8171f79 commit 4f1a9df
Show file tree
Hide file tree
Showing 4 changed files with 105 additions and 0 deletions.
1 change: 1 addition & 0 deletions core/conversion/converters/BUILD
100644 → 100755
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ cc_library(
"impl/shuffle.cpp",
"impl/softmax.cpp",
"impl/unary.cpp",
"impl/interpolate.cpp"
],
deps = [
"@tensorrt//:nvinfer",
Expand Down
34 changes: 34 additions & 0 deletions core/conversion/converters/impl/interpolate.cpp
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
#include "torch/torch.h"
#include "core/util/prelude.h"
#include "core/conversion/converters/converters.h"

namespace trtorch {
namespace core {
namespace conversion {
namespace converters {
namespace impl {
namespace {

auto interpolate_registrations = RegisterNodeConversionPatterns()
.pattern({
"aten::upsample_nearest2d(Tensor self, int[2] output_size, float? scales_h=None, float? scales_w=None) -> (Tensor)",
[](ConversionCtx* ctx, const torch::jit::Node* n, args& args) -> bool {
auto in = args[0].ITensor();

auto shape = util::toVec(in->getDimensions());

LOG_DEBUG("Shape of input is" << in);

std::cout << "TEST!" << std::endl;

return true;
}
});


} // namespace
} // namespace impl
} // namespace converters
} // namespace conversion
} // namespace core
} // namespace trtorch
70 changes: 70 additions & 0 deletions test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
# import torch.nn as nn
# import torch

# class FeatureExtractor(nn.Module):
# def __init__(self):
# super(FeatureExtractor, self).__init__()
# self.conv1 = nn.Conv2d(1, 6, 3)
# self.conv2 = nn.Conv2d(6, 16, 3)

# def forward(self, x):
# x = torch.max_pool2d(torch.relu(self.conv1(x)), (2, 2))
# x = torch.max_pool2d(torch.relu(self.conv2(x)), 2)

# return x

# class Classifier(nn.Module):
# def __init__(self):
# super(Classifier, self).__init__()

# self.fc1 = nn.Linear(16*6*6, 120)
# self.fc2 = nn.Linear(120, 84)
# self.fc3 = nn.Linear(84, 10)

# def forward(self, x):
# x = torch.flatten(x, 1)
# x = torch.relu(self.fc1(x))
# x = torch.relu(self.fc2(x))
# x = self.fc3(x)

# return x

# class LeNet(nn.Module):
# def __init__(self):
# super(LeNet, self).__init__()
# self.feat = FeatureExtractor()
# self.classifier = Classifier()

# def forward(self, x):
# x = self.feat(x)
# x = self.classifier(x)

# return x

# model = LeNet()
# model.eval()
# traced_model = torch.jit.trace(model, torch.empty([1, 1, 32, 32]))
# torch.jit.save(traced_model, 'traced_model.ts')
# torch.jit.save(torch.jit.script(model), 'script_model.ts')


import torch.nn as nn
import torch
import torch.nn.functional as F
#import trtorch

class Interp(nn.Module):
def __init__(self):
super(Interp, self).__init__()

def forward(self, x):
return F.interpolate(x, scale_factor=(5,5), mode='nearest')

model = Interp()
model.eval()
trace = torch.jit.trace(model, torch.empty([1, 1, 2, 2]))
torch.jit.save(trace, 'trace.ts')

#trtorch.check_method_op_support(trace, "forward")


Binary file added trace.ts
Binary file not shown.

0 comments on commit 4f1a9df

Please sign in to comment.