[go: nahoru, domu]
Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Optimize torch tests #3112

Merged
merged 4 commits into from
Jun 28, 2024
Merged

Optimize torch tests #3112

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
91afedf
Optimize slow torch tests
quic-mtuttle May 23, 2024
67f053d
Fix slow auto_quant tests
quic-mtuttle May 23, 2024
c1291eb
Remove redundant bn fold tests
quic-mtuttle May 23, 2024
f7cacc5
Revert bn fold test change
quic-mtuttle May 30, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions TrainingExtensions/torch/test/python/models/mnist_torch_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,15 @@ def forward(self, *inputs):
x = self.fc2(x)
return self.log_softmax(x)

class NetSmall(Net):

def __init__(self):
""" Constructor """

super(NetSmall, self).__init__()
self.fc1 = nn.Linear(7 * 7 * 64, 128)
self.fc2 = nn.Linear(128, 10)


class ExtendedNet(nn.Module):
""" Mnist Model """
Expand Down
2 changes: 1 addition & 1 deletion TrainingExtensions/torch/test/python/test_adaround_weight.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -989,7 +989,7 @@ def test_apply_adaround_using_gpu(self, dtype):
dummy_input = [x.to(dtype=dtype) for x in dummy_input]
out_before_ada = model(*dummy_input)

params = AdaroundParameters(data_loader=data_loader, num_batches=4, default_num_iterations=1000)
params = AdaroundParameters(data_loader=data_loader, num_batches=4, default_num_iterations=50)
ada_rounded_model = Adaround.apply_adaround(model, dummy_input, params, './', 'dummy')
out_after_ada = ada_rounded_model(*dummy_input)

Expand Down
21 changes: 11 additions & 10 deletions TrainingExtensions/torch/test/python/test_model_preparer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
from torchvision import models
from math import sqrt
from torch.utils.data import DataLoader
from models.test_models import SingleResidual

from aimet_common.defs import QuantScheme
from aimet_torch import elementwise_ops
Expand Down Expand Up @@ -407,9 +408,9 @@ def test_fx_with_batch_norm_folding(self):
"""
test torch fx with torchvision Resnet18 - BN fold
"""
input_shape = (1, 3, 224, 224)
input_shape = (1, 3, 32, 32)
input_tensor = torch.randn(*input_shape)
model = models.resnet18().eval()
model = SingleResidual().eval()
model_copy = copy.deepcopy(model)
folded_pairs_for_original_model = fold_all_batch_norms(model, input_shape)

Expand Down Expand Up @@ -444,9 +445,9 @@ def test_fx_with_cle(self):
"""
test torch fx with torchvision Resnet18 - Cross layer equalization
"""
input_shape = (1, 3, 224, 224)
input_shape = (1, 3, 32, 32)
input_tensor = torch.randn(*input_shape).cuda()
model = models.resnet18().cuda().eval()
model = SingleResidual().cuda().eval()
model_copy = copy.deepcopy(model)

# Perform CLE - (BN fold, ReLU6 -> ReLU replacement, CLS, HBF)
Expand Down Expand Up @@ -480,20 +481,20 @@ def test_fx_with_adaround(self):
test torch fx with torchvision Resnet18 - adaround
"""
seed_all(1)
input_shape = (1, 3, 224, 224)
input_shape = (1, 3, 32, 32)
dummy_input = torch.randn(*input_shape).cuda()
model = models.resnet18().cuda().eval()
model = SingleResidual().cuda().eval()
model_copy = copy.deepcopy(model)

# create fake data loader with image size (3, 224, 224)
data_loader = create_fake_data_loader(dataset_size=16, batch_size=16, image_size=input_shape[1:])
params = AdaroundParameters(data_loader=data_loader, num_batches=1, default_num_iterations=5)
adarounded_original_model = Adaround.apply_adaround(model, dummy_input, params, path='./',
filename_prefix='resnet18')
filename_prefix='resnet')
# Apply Adaround for transformed model
model_transformed = prepare_model(model_copy)
adarounded_transformed_model = Adaround.apply_adaround(model_transformed, dummy_input, params, path='./',
filename_prefix='resnet18')
filename_prefix='resnet')
# compare weights for very first layer
# Weights should be same
original_model_conv1_weight = adarounded_original_model.conv1.weight.clone()
Expand All @@ -514,9 +515,9 @@ def test_fx_with_bias_correction(self):
test torch fx with torchvision Resnet18 - bias correction
"""
seed_all(1)
input_shape = (1, 3, 224, 224)
input_shape = (1, 3, 32, 32)
dummy_input = torch.randn(*input_shape).cuda()
model = models.resnet18().cuda().eval()
model = SingleResidual().cuda().eval()
model_copy = copy.deepcopy(model)

# create fake data loader with image size (3, 224, 224)
Expand Down
11 changes: 6 additions & 5 deletions TrainingExtensions/torch/test/python/test_quantsim_transformers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -530,19 +530,20 @@ def test_prepare_model_with_pytorch_transformer_layer_after_act_replacement(self
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

src = torch.rand(10, 32, 512)
dummy_input = torch.rand(10, 32, 512)
src = torch.rand(1, 32, 128)
dummy_input = torch.rand(1, 32, 128)

def forward_pass(model, args):
model.eval()
with torch.no_grad():
model(dummy_input, dummy_input)

num_encoder_layers = 12
default_num_decoder_layers = 6
num_encoder_layers = 2
default_num_decoder_layers = 2

# start with a vanilla PyTorch transformer layer
transformer_model = nn.Transformer(nhead=16, num_encoder_layers=num_encoder_layers)
transformer_model = nn.Transformer(d_model=128, dim_feedforward=256, nhead=16, num_encoder_layers=num_encoder_layers,
num_decoder_layers=default_num_decoder_layers)
transformer_model.eval()

from torch import fx
Expand Down
10 changes: 6 additions & 4 deletions TrainingExtensions/torch/test/python/test_weight_pad_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@
from aimet_torch.tensor_quantizer import StaticGridPerTensorQuantizer, StaticGridPerChannelQuantizer
from aimet_torch.weight_padding_utils import recompute_scale, recompute_encodings, weight_pad, WeightPaddingParams

from models.test_models import TinyModel


def evaluate(model: torch.nn.Module, dummy_input: torch.Tensor):
"""
Expand Down Expand Up @@ -186,8 +188,8 @@ def test_weight_pad_in_place_per_channel(self):
assert val % 16 == 0

def test_weight_pad_in_place_per_tensor(self):
model = models.resnet50(pretrained=True)
dummy_input = torch.randn(1, 3, 224, 224)
model = TinyModel()
dummy_input = torch.randn(1, 3, 32, 32)
sim = QuantizationSimModel(model, dummy_input, quant_scheme=QuantScheme.post_training_tf_enhanced,
default_param_bw=16, default_output_bw=16)
sim.compute_encodings(evaluate, dummy_input)
Expand Down Expand Up @@ -215,8 +217,8 @@ def test_weight_pad_in_place_per_tensor(self):
assert val % 16 == 0

def test_weight_pad_export_per_tensor(self):
model = models.resnet50(pretrained=True)
dummy_input = torch.randn(1, 3, 224, 224)
model = TinyModel()
dummy_input = torch.randn(1, 3, 32, 32)
sim = QuantizationSimModel(model, dummy_input, quant_scheme=QuantScheme.post_training_tf_enhanced,
default_param_bw=16, default_output_bw=16)
sim.compute_encodings(evaluate, dummy_input)
Expand Down
4 changes: 2 additions & 2 deletions TrainingExtensions/torch/test/python/test_weight_svd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -874,7 +874,7 @@ def test_prune_layer(self):

def test_prune_model_2_layers(self):

model = mnist_model.Net()
model = mnist_model.NetSmall()

# Create a layer database
input_shape = (1, 1, 28, 28)
Expand All @@ -893,7 +893,7 @@ def test_prune_model_2_layers(self):
fc1_b = layer_db.find_layer_by_name('fc1.1')

self.assertEqual(3136, fc1_a.module.in_features)
self.assertEqual(1024, fc1_b.module.out_features)
self.assertEqual(128, fc1_b.module.out_features)

conv2_a = layer_db.find_layer_by_name('conv2.0')
conv2_b = layer_db.find_layer_by_name('conv2.1')
Expand Down
2 changes: 1 addition & 1 deletion TrainingExtensions/torch/test/python/v2/test_adaround.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -982,7 +982,7 @@ def test_apply_adaround_using_gpu(self, dtype):
dummy_input = [x.to(dtype=dtype) for x in dummy_input]
out_before_ada = model(*dummy_input)

params = AdaroundParameters(data_loader=data_loader, num_batches=4, default_num_iterations=1000)
params = AdaroundParameters(data_loader=data_loader, num_batches=4, default_num_iterations=50)
ada_rounded_model = Adaround.apply_adaround(model, dummy_input, params, './', 'dummy')
out_after_ada = ada_rounded_model(*dummy_input)

Expand Down
Loading