summary | shortlog | log | commit | commitdiff | tree
raw | patch | inline | side by side (from parent 1: 2ee526e)
raw | patch | inline | side by side (from parent 1: 2ee526e)
author | Manu Mathew <a0393608@ti.com> | |
Tue, 18 Feb 2020 18:37:36 +0000 (00:07 +0530) | ||
committer | Manu Mathew <a0393608@ti.com> | |
Tue, 18 Feb 2020 18:37:36 +0000 (00:07 +0530) |
index 64f2f8ed157166b59acded0b1e7a90f90be9a738..be09729ac1021d2782141fdac90e8205893edb7a 100644 (file)
import os
import torch
import os
import torch
-import torchvision
import datetime
import datetime
+import torchvision as vision
+# from pytorch_jacinto_ai import vision
# dependencies
# Anaconda Python 3.7 for Linux - download and install from: https://www.anaconda.com/distribution/
# dependencies
# Anaconda Python 3.7 for Linux - download and install from: https://www.anaconda.com/distribution/
# some parameters - modify as required
date = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
dataset_name = 'image_folder_classification'
# some parameters - modify as required
date = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
dataset_name = 'image_folder_classification'
-model_name = 'resnet50'
+model_names = ['mobilenet_v2', 'resnet18', 'resnet50', 'resnext50_32x4d', 'shufflenet_v2_x1_0']
img_resize = (256,256)
rand_crop = (224,224)
# the saving path - you can choose any path
save_path = './data/checkpoints'
img_resize = (256,256)
rand_crop = (224,224)
# the saving path - you can choose any path
save_path = './data/checkpoints'
-save_path = os.path.join(save_path, dataset_name, date + '_' + dataset_name + '_' + model_name)
+save_path = os.path.join(save_path, dataset_name, date + '_' + dataset_name)
save_path += '_resize{}x{}_traincrop{}x{}'.format(img_resize[1], img_resize[0], rand_crop[1], rand_crop[0])
os.makedirs(save_path, exist_ok=True)
save_path += '_resize{}x{}_traincrop{}x{}'.format(img_resize[1], img_resize[0], rand_crop[1], rand_crop[0])
os.makedirs(save_path, exist_ok=True)
-# create the model - replace with your model
-model = torchvision.models.resnet50(pretrained=True)
-
# create a rand input
rand_input = torch.rand(1, 3, rand_crop[0], rand_crop[1])
# create a rand input
rand_input = torch.rand(1, 3, rand_crop[0], rand_crop[1])
-# write the onnx model
-torch.onnx.export(model, rand_input, os.path.join(save_path, 'model.onnx'), export_params=True, verbose=False)
+for model_name in model_names:
+ # create the model - replace with your model
+ model = vision.models.__dict__[model_name](pretrained=True)
+ model.eval()
+
+ # write pytorch model
+ model_path=os.path.join(save_path, f'{model_name}_model.pth')
+ traced_model = torch.jit.trace(model, rand_input)
+ torch.jit.save(traced_model, model_path)
+
+ # write pytorch sate dict
+ model_path=os.path.join(save_path, f'{model_name}_state_dict.pth')
+ torch.save(model.state_dict(), model_path)
+
+ # write the onnx model
+ opset_version=9
+ model_path=os.path.join(save_path, f'{model_name}_opset{opset_version}.onnx')
+ torch.onnx.export(model, rand_input, model_path, export_params=True, verbose=False, opset_version=opset_version)
diff --git a/modules/pytorch_jacinto_ai/vision/models/resnet.py b/modules/pytorch_jacinto_ai/vision/models/resnet.py
index 788fd8b46a51a47ea97cbd38734feb78f68e5457..f4f4026fffb3d47d93036696a633214cf730e17f 100644 (file)
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
def _resnet(arch, block, layers, pretrained, progress, **kwargs):
model = ResNet(block, layers, **kwargs)
- if pretrained:
+ if pretrained is True:
+ change_names_dict = kwargs.get('change_names_dict', None)
+ state_dict = load_state_dict_from_url(model_urls[arch],
+ progress=progress)
+ model.load_weights(state_dict, change_names_dict=change_names_dict)
+ elif pretrained:
change_names_dict = kwargs.get('change_names_dict', None)
download_root = kwargs.get('download_root', None)
model.load_weights(pretrained, change_names_dict=change_names_dict, download_root=download_root)
change_names_dict = kwargs.get('change_names_dict', None)
download_root = kwargs.get('download_root', None)
model.load_weights(pretrained, change_names_dict=change_names_dict, download_root=download_root)
diff --git a/modules/pytorch_jacinto_ai/vision/models/shufflenetv2.py b/modules/pytorch_jacinto_ai/vision/models/shufflenetv2.py
index 806528bea8ef82f0472b5ecedd6c6b00567d25db..be529d6a775bcfef4fbee7ebc62933de95c12a5b 100644 (file)
+from collections import OrderedDict
import torch
import torch.nn as nn
from .utils import load_state_dict_from_url
import torch
import torch.nn as nn
from .utils import load_state_dict_from_url
nn.ReLU(inplace=True),
)
layers += [('conv5',conv5)]
nn.ReLU(inplace=True),
)
layers += [('conv5',conv5)]
- self.features = torch.nn.Sequential(layers)
+ self.features = torch.nn.Sequential(OrderedDict(layers))
if self.num_classes is not None:
self.classifier = nn.Linear(output_channels, num_classes)
if self.num_classes is not None:
self.classifier = nn.Linear(output_channels, num_classes)