def test_manual_mode(self):
torch.cuda.empty_cache()
net = mnist_model.Net()
model = net.to(torch.device('cuda'))
# Adding wrapper to first convolution layer
for module_name, module_ref in model.named_children():
if module_name is 'conv1':
quantized_module = QcPostTrainingWrapper(
module_ref,
weight_bw=8,
activation_bw=8,
round_mode='nearest',
quant_scheme=QuantScheme.post_training_tf)
setattr(model, module_name, quantized_module)
sim = QuantizationSimModel(model,
dummy_input=torch.rand(1, 1, 28, 28).cuda())
# Quantize the untrained MNIST model
sim.compute_encodings(self.forward_pass, forward_pass_callback_args=5)
# Run some inferences
mnist_torch_model.evaluate(model=sim.model,
iterations=100,
use_cuda=True)
# train the model again
mnist_model.train(model=sim.model,
epochs=1,
num_batches=3,
batch_callback=check_if_layer_weights_are_updating,
use_cuda=True)
def train_model(self, model, layer, train_flag=True):
"""
Trains a model
:param model: Model to be trained
:param layer: layer which has to be fine tuned
:param train_flag: Default: True. If ture the model gets trained
:return:
"""
if train_flag:
mnist_torch_model.train(model, epochs=1, use_cuda=True, batch_size=50, batch_callback=None)
self._layer_db.append(layer)
def test_retraining_on_quantized_model_second_step(self):
torch.cuda.empty_cache()
sim = load_checkpoint(os.path.join(path, 'checkpoint.pt'))
# re-train the model for entire one epoch
mnist_model.train(model=sim.model,
epochs=1,
num_batches=3,
batch_callback=check_if_layer_weights_are_updating,
use_cuda=True)
def test_pretrained_mnist_quantize(self):
model = mnist_model.Net().to('cpu')
mnist_model.train(model, epochs=1, num_batches=10)
torch.save(model, 'data/mnist.pth')
pretrained_model = torch.load('data/mnist.pth')
quantizer = q.Quantizer(model=pretrained_model, use_cuda=False)
# Quantize
quantizer.quantize_net(bw_params=8,
bw_acts=8,
run_model=mnist_model.evaluate,
iterations=10)
def test_retraining_on_quantized_model_first_step(self):
torch.cuda.empty_cache()
model = mnist_model.Net().to(torch.device('cuda'))
sim = QuantizationSimModel(model,
default_output_bw=4,
default_param_bw=4,
dummy_input=torch.rand(1, 1, 28, 28).cuda())
# Quantize the untrained MNIST model
sim.compute_encodings(self.forward_pass, forward_pass_callback_args=5)
# train the model for entire one epoch
mnist_model.train(model=sim.model,
epochs=1,
num_batches=3,
batch_callback=check_if_layer_weights_are_updating,
use_cuda=True)
# Checkpoint the model
save_checkpoint(sim, os.path.join(path, 'checkpoint.pt'))
def test_with_finetuning(self):
torch.cuda.empty_cache()
model = mnist_model.Net().to(torch.device('cuda'))
mnist_torch_model.evaluate(model=model, iterations=None, use_cuda=True)
sim = QuantizationSimModel(model,
dummy_input=torch.rand(1, 1, 28, 28).cuda())
# Quantize the untrained MNIST model
sim.compute_encodings(self.forward_pass, forward_pass_callback_args=5)
# Run some inferences
mnist_torch_model.evaluate(model=sim.model,
iterations=None,
use_cuda=True)
# train the model again
mnist_model.train(sim.model,
epochs=1,
num_batches=3,
batch_callback=check_if_layer_weights_are_updating,
use_cuda=True)
cpu_output_files = os.path.join('./', 'data', 'mnist_trained_on_CPU.pth')
gpu_output_files = os.path.join('./', 'data', 'mnist_trained_on_GPU.pth')
if os.path.isfile(cpu_output_files) or os.path.isfile(gpu_output_files):
logger.info('Mnist model .pth generation not needed')
else:
torch.manual_seed(1)
torch.backends.cudnn.deterministic = True
if use_cuda:
model = mnist_torch_model.Net().to("cuda")
else:
model = mnist_torch_model.Net().to("cpu")
mnist_torch_model.train(model,
epochs=1,
use_cuda=use_cuda,
batch_size=50,
batch_callback=None)
# create directory
if not os.path.isdir('./data'):
os.mkdir('./data')
if use_cuda:
torch.save(model, gpu_output_files)
model = model.to("cpu")
torch.save(model, cpu_output_files)
# ##############################################################################################