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) # ##############################################################################################