<jupyter_start><jupyter_text>Create data iterators<jupyter_code>batch_size = 64
train_folder, val_folder = get_image_folders()

train_iterator = DataLoader(
    train_folder, batch_size=batch_size, num_workers=4,
    shuffle=True, pin_memory=True
)

val_iterator = DataLoader(
    val_folder, batch_size=256, num_workers=4,
    shuffle=False, pin_memory=True
)

# number of training samples
train_size = len(train_folder.imgs)
train_size<jupyter_output><empty_output><jupyter_text>Model<jupyter_code>from get_densenet import get_model
model, loss, optimizer = get_model(learning_rate=LEARNING_RATE)

# load pretrained model, accuracy ~73%
model.load_state_dict(torch.load('../vanilla_densenet_big/model_step5.pytorch_state'))<jupyter_output><empty_output><jupyter_text>keep copy of full precision kernels<jupyter_code># copy almost all full precision kernels of the model
all_fp_kernels = [
    Variable(kernel.data.clone(), requires_grad=True) 
    for kernel in optimizer.param_groups[1]['params']
]
# all_fp_kernels - kernel tensors of all convolutional layers 
# (with the exception of the first conv layer)<jupyter_output><empty_output><jupyter_text>initial quantization<jupyter_code># scaling factors for each quantized layer
initial_scaling_factors = []
# these kernels will be quantized
all_kernels = [kernel for kernel in optimizer.param_groups[1]['params']]
for k, k_fp in zip(all_kernels, all_fp_kernels):
    
    # choose initial scaling factors 
    w_p_initial, w_n_initial = initial_scales(k_fp.data)
    initial_scaling_factors += [(w_p_initial, w_n_initial)]
    
    # do quantization
    k.data = quantize(k_fp.data, w_p_initial, w_n_initial, t=HYPERPARAMETER_T)<jupyter_output><empty_output><jupyter_text>parameter updaters<jupyter_code># optimizer for updating only all_fp_kernels
optimizer_fp = optim.Adam(all_fp_kernels, lr=LEARNING_RATE)
# optimizer for updating only scaling factors
optimizer_sf = optim.Adam([
    Variable(torch.FloatTensor([w_p, w_n]).cuda(), requires_grad=True) 
    for w_p, w_n in initial_scaling_factors
], lr=LEARNING_RATE)<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>from torch.optim.lr_scheduler import ReduceLROnPlateau

class lr_scheduler_list:
    """ReduceLROnPlateau for a list of optimizers."""
    def __init__(self, optimizer_list):
        self.lr_scheduler_list = [
            ReduceLROnPlateau(
                optimizer, mode='max', factor=0.1, patience=3, 
                verbose=True, threshold=0.01, threshold_mode='abs'
            ) 
            for optimizer in optimizer_list
        ]
    
    def step(self, test_accuracy):
        for scheduler in self.lr_scheduler_list:
            scheduler.step(test_accuracy)

n_epochs = 15
n_batches = ceil(train_size/batch_size)

# total number of batches in the train set
n_batches
%%time
optimizer_list = [optimizer, optimizer_fp, optimizer_sf]

def optimization_step_fn(model, loss, x_batch, y_batch):
    return optimization_step(
        model, loss, x_batch, y_batch, 
        optimizer_list=optimizer_list,
        t=HYPERPARAMETER_T
    )
all_losses = train(
    model, loss, optimization_step_fn,
    train_iterator, val_iterator, n_epochs,
    lr_scheduler=lr_scheduler_list(optimizer_list)        
)
# epoch logloss  accuracy    top5_accuracy time  (first value: train, second value: val)
# backup
model.cpu();
torch.save(model.state_dict(), 'model_ternary_quantization.pytorch_state')<jupyter_output><empty_output><jupyter_text>Continue training<jupyter_code># reduce learning rate
for optimizer in optimizer_list:
    for group in optimizer.param_groups:
        group['lr'] = 1e-5
n_epochs = 5
model.cuda();
%%time
def optimization_step_fn(model, loss, x_batch, y_batch):
    return optimization_step(
        model, loss, x_batch, y_batch, 
        optimizer_list=optimizer_list,
        t=HYPERPARAMETER_T
    )
all_losses = train(
    model, loss, optimization_step_fn,
    train_iterator, val_iterator, n_epochs       
)
# epoch logloss  accuracy    top5_accuracy time  (first value: train, second value: val)<jupyter_output>0  1.312 1.362  0.660 0.654  0.875 0.866  649.648
1  1.243 1.341  0.675 0.661  0.885 0.868  651.109
2  1.222 1.366  0.678 0.657  0.889 0.866  651.324
3  1.207 1.334  0.682 0.664  0.890 0.870  651.090
4  1.195 1.350  0.685 0.659  0.891 0.869  651.020
CPU times: user 57min 55s, sys: 8min 45s, total: 1h 6min 41s
Wall time: 54min 14s<jupyter_text>Final save<jupyter_code>model.cpu();
torch.save(model.state_dict(), 'model_ternary_quantization.pytorch_state')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09314047775586237, 'fraction_numerical': 0.042735042735042736, 'mean_word_length': 4.662531017369727, 'pattern_counts': {'":': 0, '<': 32, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '149612', 'license': 'MIT', 'n_tokens_mistral': 1764, 'n_tokens_neox': 1624, 'n_words': 343, 'path': 'ttq_densenet_big/train.ipynb', 'repo_name': 'dan-fritchman/trained-ternary-quantization', 'size': '4563'}