Pull fixes from develop in examples.

Author: jmohdyusof

examples/histogen/extract_code_default_model.txt

@@ -3,5 +3,8 @@ size = 256
 batch_size = 128
 use_gpus = True
 ckpt_directory = './'
+ckpt_restart = 'checkpoint/vqvae_001.pt'
+lmdb_filename = 'lmdb_001'
+data_dir = '../../Data/Examples/histogen/svs_pngs/'
 
 

examples/histogen/sample_default_model.txt

@@ -8,5 +8,4 @@ batch_size = 8
 use_gpus = True
 
 [Monitor_Params]
-timeout=3600
-
+timeout = 3600

examples/histogen/train_pixelsnail_baseline_pytorch.py

@@ -0,0 +1,252 @@
+import sys
+import os
+
+import numpy as np
+import torch
+from torch import nn, optim
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from argparse import SUPPRESS
+
+try:
+    from apex import amp
+
+except ImportError:
+    amp = None
+
+from dataset import LMDBDataset
+from pixelsnail import PixelSNAIL
+from scheduler import CycleScheduler
+
+file_path = os.path.dirname(os.path.realpath(__file__))
+lib_path = os.path.abspath(os.path.join(file_path, '..'))
+sys.path.append(lib_path)
+lib_path2 = os.path.abspath(os.path.join(file_path, '..', '..', 'common'))
+sys.path.append(lib_path2)
+
+
+import candle
+
+additional_definitions = [
+    {'name': 'sched_mode',
+        'type': str,
+        'default': None,
+        'help': 'Mode of learning rate scheduler'},
+    {'name': 'lmdb_filename',
+        'type': str,
+        'default': SUPPRESS,
+        'help': 'lmdb dataset path'},
+    {'name': 'amp',
+        'type': str,
+        'default': 'O0',
+        'help': ''},
+    {'name': 'hier',
+        'type': str,
+        'default': 'top',
+        'help': ''},
+    {'name': 'channel',
+        'type': int,
+        'default': 256,
+        'help': ''},
+    {'name': 'n_res_block',
+        'type': int,
+        'default': 4,
+        'help': ''},
+    {'name': 'n_res_channel',
+        'type': int,
+        'default': 256,
+        'help': ''},
+    {'name': 'n_out_res_block',
+        'type': int,
+        'default': 0,
+        'help': ''},
+    {'name': 'n_cond_res_block',
+        'type': int,
+        'default': 3,
+        'help': ''},
+    {'name': 'ckpt_restart',
+        'type': str,
+        'default': None,
+        'help': 'Checkpoint to restart from'},
+]
+
+required = [
+    'batch_size',
+    'epochs',
+    'hier',
+    'learning_rate',
+    'channel',
+    'n_res_block',
+    'n_res_channel',
+    'n_out_res_block',
+    'n_cond_res_block',
+    'dropout',
+    'amp',
+    'sched_mode',
+    'lmdb_filename',
+]
+
+
+class TrPxSnBk(candle.Benchmark):
+
+    def set_locals(self):
+        """Functionality to set variables specific for the benchmark
+        - required: set of required parameters for the benchmark.
+        - additional_definitions: list of dictionaries describing the additional parameters for the
+        benchmark.
+        """
+
+        if required is not None:
+            self.required = set(required)
+        if additional_definitions is not None:
+            self.additional_definitions = additional_definitions
+
+
+def initialize_parameters(default_model='train_pixelsnail_default_model.txt'):
+
+    # Build benchmark object
+    trpsn = TrPxSnBk(file_path, default_model, 'pytorch',
+                     prog='train_pixelsnail_baseline',
+                     desc='Histology train pixelsnail - Examples')
+
+    print("Created sample benchmark")
+
+    # Initialize parameters
+    gParameters = candle.finalize_parameters(trpsn)
+    print("Parameters initialized")
+
+    return gParameters
+
+
+def train(args, epoch, loader, model, optimizer, scheduler, device):
+    loader = tqdm(loader)
+
+    criterion = nn.CrossEntropyLoss()
+
+    for i, (top, bottom, label) in enumerate(loader):
+        model.zero_grad()
+
+        top = top.to(device)
+
+        if args.hier == 'top':
+            target = top
+            out, _ = model(top)
+
+        elif args.hier == 'bottom':
+            bottom = bottom.to(device)
+            target = bottom
+            out, _ = model(bottom, condition=top)
+
+        loss = criterion(out, target)
+        loss.backward()
+
+        if scheduler is not None:
+            scheduler.step()
+        optimizer.step()
+
+        _, pred = out.max(1)
+        correct = (pred == target).float()
+        accuracy = correct.sum() / target.numel()
+
+        lr = optimizer.param_groups[0]['lr']
+
+        loader.set_description(
+            (
+                f'epoch: {epoch + 1}; loss: {loss.item():.5f}; '
+                f'acc: {accuracy:.5f}; lr: {lr:.5f}'
+            )
+        )
+
+
+class PixelTransform:
+    def __init__(self):
+        pass
+
+    def __call__(self, input):
+        ar = np.array(input)
+
+        return torch.from_numpy(ar).long()
+
+
+def run(params):
+
+    args = candle.ArgumentStruct(**params)
+    # Configure GPUs
+    ndevices = torch.cuda.device_count()
+    if ndevices < 1:
+        raise Exception('No CUDA gpus available')
+
+    device = 'cuda'
+
+    dataset = LMDBDataset(args.lmdb_filename)
+    loader = DataLoader(
+        dataset, batch_size=args.batch_size, shuffle=True, num_workers=4, drop_last=True
+    )
+
+    ckpt = {}
+
+    if args.ckpt_restart is not None:
+        ckpt = torch.load(args.ckpt_restart)
+        args = ckpt['args']
+
+    if args.hier == 'top':
+        model = PixelSNAIL(
+            [32, 32],
+            512,
+            args.channel,
+            5,
+            4,
+            args.n_res_block,
+            args.n_res_channel,
+            dropout=args.dropout,
+            n_out_res_block=args.n_out_res_block,
+        )
+
+    elif args.hier == 'bottom':
+        model = PixelSNAIL(
+            [64, 64],
+            512,
+            args.channel,
+            5,
+            4,
+            args.n_res_block,
+            args.n_res_channel,
+            attention=False,
+            dropout=args.dropout,
+            n_cond_res_block=args.n_cond_res_block,
+            cond_res_channel=args.n_res_channel,
+        )
+
+    if 'model' in ckpt:
+        model.load_state_dict(ckpt['model'])
+
+    model = model.to(device)
+    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate)
+
+    if amp is not None:
+        model, optimizer = amp.initialize(model, optimizer, opt_level=args.amp)
+
+    model = nn.DataParallel(model)
+    model = model.to(device)
+
+    scheduler = None
+    if args.sched_mode == 'cycle':
+        scheduler = CycleScheduler(
+            optimizer, args.learning_rate, n_iter=len(loader) * args.epochs, momentum=None
+        )
+
+    for i in range(args.epochs):
+        train(args, i, loader, model, optimizer, scheduler, device)
+        torch.save(
+            {'model': model.module.state_dict(), 'args': args},
+            f'{args.ckpt_directory}/checkpoint/pixelsnail_{args.hier}_{str(i + 1).zfill(3)}.pt',
+        )
+
+
+def main():
+    params = initialize_parameters()
+    run(params)
+
+
+if __name__ == '__main__':
+    main()

examples/histogen/train_pixelsnail_default_model.txt

@@ -1,4 +1,5 @@
 [Global_Params]
+lmdb_filename = 'lmdb_001'
 epochs = 420
 batch_size = 32
 learning_rate = 3e-4
@@ -12,5 +13,3 @@ dropout = 0.1
 amp = 'O0'
 use_gpus = True
 ckpt_directory = './'
-
-

examples/histogen/train_vqvae_baseline_pytorch.py

@@ -43,7 +43,7 @@
         'type': str,
         'default': SUPPRESS,
         'help': 'dataset path'},
-    {'name': 'size',
+    {'name': 'image_size',
         'type': int,
         'default': 256,
         'help': 'Image size to use'},
@@ -55,7 +55,7 @@
     'epochs',
     'learning_rate',
     'sched_mode',
-    'size',
+    'image_size',
 ]
 
 
@@ -168,8 +168,8 @@ def config_and_train(args):
 
     transform = transforms.Compose(
         [
-            transforms.Resize(args.size),
-            transforms.CenterCrop(args.size),
+            transforms.Resize(args.image_size),
+            transforms.CenterCrop(args.image_size),
             transforms.ToTensor(),
             transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),
         ]
@@ -212,7 +212,7 @@ def fetch_data(params):
     if params['data_dir'] is None:
         params['data_dir'] = candle.fetch_file(data_url + params['train_data'], subdir='Examples/histogen')
     else:
-        tempfile = candle.fetch_file(data_url + params['train_data'], cache_subdir='Examples/histogen')
+        tempfile = candle.fetch_file(data_url + params['train_data'], subdir='Examples/histogen')
         params['data_dir'] = os.path.join(os.path.dirname(tempfile), params['data_dir'])
 
 

examples/histogen/train_vqvae_default_model.txt

@@ -5,7 +5,6 @@ data_dir = 'svs_pngs'
 epochs = 560
 learning_rate = 3e-4
 batch_size = 128
-use_gpus = True
+n_gpu_per_machine = 1
 ckpt_directory = './'
-
-
+image_size = 256

examples/rnngen/infer_rnngen_baseline_pytorch.py

@@ -10,7 +10,7 @@
 
 from model.model import CharRNN
 from model.vocab import START_CHAR, END_CHAR
-from train import get_vocab_from_file
+from model.vocab import get_vocab_from_file
 
 file_path = os.path.dirname(os.path.realpath(__file__))
 lib_path = os.path.abspath(os.path.join(file_path, '..'))
@@ -70,6 +70,7 @@
     'output',
     'input',
     'nsamples',
+    'model',
 ]
 
 
@@ -175,6 +176,18 @@ def run(params):
     print("Note: This script is very picky. Please check device output to see where this is running. ")
     args = candle.ArgumentStruct(**params)
 
+    data_url = args.data_url
+
+    if args.model == 'ft_goodperforming_model.pt':
+        file = 'pilot1/ft_goodperforming_model.pt'
+    elif args.model == 'ft_poorperforming_model.pt':
+        file = 'pilot1/ft_poorperforming_model.pt'
+    else: # Corresponding to args.model == 'autosave.model.pt':
+        file = 'mosesrun/autosave.model.pt'
+
+    print('Recovering trained model')
+    trained = candle.fetch_file(data_url + file, subdir='examples/rnngen')
+
     # Configure GPU
     if args.use_gpus and torch.cuda.is_available():
         device = 'cuda'
@@ -188,10 +201,12 @@ def run(params):
     model = CharRNN(len(vocab), len(vocab), max_len=args.maxlen).to(device)
     optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)
 
-    pt = torch.load(args.logdir + "/" + args.model, map_location=device)
+    print("Loading trained model.")
+    pt = torch.load(trained, map_location=device)
     model.load_state_dict(pt['state_dict'])
     optimizer.load_state_dict(pt['optim_state_dict'])
 
+    print("Applying to loaded data")
     total_sampled = 0
     total_valid = 0
     total_unqiue = 0

examples/rnngen/infer_rnngen_default_model.txt

@@ -1,4 +1,5 @@
 [Global_Params]
+data_url = 'ftp://ftp.mcs.anl.gov/pub/candle/public/benchmarks/Examples/rnngen/'
 input = 'mosesrun/'
 logdir = 'mosesrun/'
 output = 'samples.txt'