<jupyter_start><jupyter_text>Update network architecture and try different style image  Import<jupyter_code>import os
import sys
import time
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.transforms.functional as F
from IPython.core.debugger import set_trace
from PIL import Image
sys.path.append('../../')
from style_transfer.data import *
from style_transfer.models import *<jupyter_output><empty_output><jupyter_text>Config<jupyter_code>dir_data = Path('../../data')<jupyter_output><empty_output><jupyter_text>Utility<jupyter_code>def save_checkpoint(model, optimizer, file_checkpoint):
    file_checkpoint.parent.mkdir(exist_ok=True)
    torch.save({'model': model.state_dict(),
                'optimizer': optimizer.state_dict()},
               file_checkpoint)
def load_checkpoint(model, optimizer, file_checkpoint):
    checkpoint = torch.load(file_checkpoint)
    model.load_state_dict(checkpoint['model'])
    optimizer.load_state_dict(checkpoint['optimizer'])
    return model, optimizer
def epoch(model, dl, loss, optimizer, device, mode):
    metrics_e, total_e = dict.fromkeys(loss.get_metric_names(), 0), 0
    model.train() if mode == 'trn' else model.eval()
    with torch.set_grad_enabled(mode == 'trn'):
        for X, y in dl:
            X, y = X.to(device), y.to(device)

            # Get loss and batch metrics
            if mode == 'trn': optimizer.zero_grad()            
            y_hat = model(X)
            l, metrics_b = loss(y_hat, y)            
            if mode == 'trn': l.backward(), optimizer.step()

            # Scale metrics
            for key in metrics_e.keys(): metrics_e[key] += metrics_b[key]*X.shape[0]
            total_e += X.shape[0] # Track total number of items in epoch

    # Get mean of metrics
    for key in metrics_e.keys(): metrics_e[key] /= total_e
    return metrics_e<jupyter_output><empty_output><jupyter_text>Train Config<jupyter_code># Config
num_workers = 4
pin_memory = True
dir_styl = dir_data/'style'
dir_cont = dir_data/'content'

# Device
device = torch.device('cuda')

# Model
in_channels = 3
layout_encoder = [[  32, 7, 1],
                  [  64, 5, 1],
                  [ 128, 3, 1],
                  [ 256, 3, 8]]
encoder = get_encoder(in_channels, layout_encoder)
model = AutoEncoder(encoder)
model = model.to(device)

# Optimizer
optimizer = optim.Adam(model.parameters())<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>def _get_dl_trn(sz, batch_size, num_samples):
    tfms = transforms.Compose([transforms.ToTensor(),
                               transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                                                    std=[0.229, 0.224, 0.225])])

    return SubsetDataloader(ImgCSVDataset(dir_cont/f'formatted_{sz}'/'trn'/'imgs',
                                          dir_cont/'trn.csv', 
                                          tfms=tfms), 
                            batch_size=batch_size,
                            num_workers=num_workers,
                            pin_memory=pin_memory,
                            num_samples=num_samples)
def _set_lr(lr): 
    for p in optimizer.param_groups: 
        p['lr'] = lr
def _fit(num_epochs, dl_trn, loss, file_checkpoint):    
    %matplotlib notebook

    fig, axs = plt.subplots(2, figsize=(8,8))
    for num_epoch in range(num_epochs):    
        # Train
        metrics_trn = epoch(model, dl_trn, loss, optimizer, device, 'trn')

        # Print metrics
        print(pd.DataFrame(metrics_trn, index=[num_epoch]).to_string(header=num_epoch == 0,
                                                                     float_format=lambda x: f'{x:14.2f}',
                                                                     col_space=8))
        # Checkpoint
        save_checkpoint(model, optimizer, file_checkpoint)

        # Plot Images
        model.eval()
        with torch.no_grad():
            y_hat = model(img_test)
            axs[0].imshow(denorm(torch2np(img_test[0])))
            axs[1].imshow(denorm(torch2np(y_hat[0])))

        # Update plot
        fig.canvas.draw()

    %matplotlib inline
def _train(num_epochs, lr, sz, batch_size, num_samples, loss, name):
    _set_lr(lr)
    dl_trn = _get_dl_trn(sz, batch_size, num_samples)
    _fit(num_epochs, dl_trn, loss, Path(f'checkpoints/{name}.pth'))<jupyter_output><empty_output><jupyter_text>Stage 1<jupyter_code>lr = 1e-5
sz = 256
batch_size =  4
num_samples= 10000    

img_test = Image.open(dir_cont/'artists'/'layne.jpg')
img_test = F.center_crop(F.resize(img_test, 128), 128)
img_test = img2tensor(img_test)

loss = FeatureLoss(models.vgg16(pretrained=True).features, 
                   F.center_crop(F.resize(Image.open(dir_styl/'dora.jpg'), 128), 128), 
                   sz, 
                   [0,1,2,3], 
                   [2],
                   [1e8, 1e8, 1e8, 1e8],
                   [1e3])

_train(100, lr, sz, batch_size, num_samples, loss, 'stage1')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11402836335760828, 'fraction_numerical': 0.020697585281717133, 'mean_word_length': 2.456291390728477, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 25, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 12, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '528889', 'license': 'Apache-2.0', 'n_tokens_mistral': 1771, 'n_tokens_neox': 1680, 'n_words': 341, 'path': 'training/attempt_4/train.ipynb', 'repo_name': 'justinblaber/style_transfer', 'size': '5218'}