<jupyter_start><jupyter_text>Train model<jupyter_code>class p:
    '''Parameters for Gaussian mixture simulation
    '''
    # parameters for generating data
    seed = 13
    dataset = "dsprites"
    
    # parameters for model architecture
    model_type = "Burgess"
    latent_dim = 10 
    img_size = (1, 64, 64)
    
    # parameters for training
    train_batch_size = 64
    test_batch_size = 100
    lr = 1e-4
    rec_dist = "bernoulli"
    reg_anneal = 0
    num_epochs = 100
    
    # hyperparameters for loss
    beta = 0.01
    lamPT = 0.0
    lamNN = 0.0
    lamH = 0.0
    lamSP = 0.0
    
    # parameters for exp
    warm_start = None # which parameter to warm start with respect to
    seq_init = 1      # value of warm_start parameter to start with respect to
    
    # SAVE MODEL
    out_dir = "/home/ubuntu/local-vae/notebooks/ex_dsprites/results" # wooseok's setup
#     out_dir = '/scratch/users/vision/chandan/local-vae' # chandan's setup
    dirname = "vary"
    pid = ''.join(["%s" % randint(0, 9) for num in range(0, 10)])

    def _str(self):
        vals = vars(p)
        return 'beta=' + str(vals['beta']) + '_lamPT=' + str(vals['lamPT']) + '_lamNN=' + str(vals['lamNN']) + '_lamSP=' + str(vals['lamSP']) \
                + '_seed=' + str(vals['seed']) + '_pid=' + vals['pid']
    
    def _dict(self):
        return {attr: val for (attr, val) in vars(self).items()
                 if not attr.startswith('_')}
# create dir
out_dir = opj(p.out_dir, p.dirname)
os.makedirs(out_dir, exist_ok=True)  

# seed
random.seed(p.seed)
np.random.seed(p.seed)
torch.manual_seed(p.seed)

# get dataloaders
train_loader = get_dataloaders(p.dataset,
                               batch_size=p.train_batch_size,
                               logger=None)

# prepare model
model = init_specific_model(model_type=p.model_type, 
                            img_size=p.img_size,
                            latent_dim=p.latent_dim,
                            hidden_dim=None).to(device)
# model.load_state_dict(torch.load('results/dsprites.pth'))

# train
optimizer = torch.optim.Adam(model.parameters(), lr=p.lr)
loss_f = get_loss_f(decoder=model.decoder, **vars(p))
trainer = Trainer(model, optimizer, loss_f, device=device)
trainer(train_loader, epochs=p.num_epochs)
im = iter(train_loader).next()[0].to(device)
recon_data, latent_dist, latent_sample = model(im)
viz_im_r(im[10,0], recon_data[10,0])
torch.save(model.state_dict(), 'results/dsprites.pth')
from scipy import stats
from visualize import Visualizer
import torch.nn.functional as F
from visualize import make_grid_img, get_samples
from torchvision.utils import make_grid
model_dir = opj('results/cosmo')
model.eval()  # don't sample from latent: use mean
viz = Visualizer(model=model,
                 model_dir=model_dir,
                 save_images=False)
n_rows = 10
n_cols = 10
size = (n_rows, n_cols)
num_samples = n_rows * n_cols
samples = get_samples(train_loader.dataset, num_samples)
a = viz.traversals(data=samples[0:1, ...],
               n_per_latent=n_cols,
               n_latents=n_rows,
               is_force_return=True)
plt.figure(figsize=(10,10))
plt.imshow(a, cmap='magma')
plt.axis('off')
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12472985489348565, 'fraction_numerical': 0.01667181228774313, 'mean_word_length': 2.884892086330935, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '87848', 'license': 'MIT', 'n_tokens_mistral': 1201, 'n_tokens_neox': 1132, 'n_words': 237, 'path': 'notebooks/ex_dsprites/train_disvae.ipynb', 'repo_name': 'csinva/local-vae', 'size': '3239'}