test_datasets = [
get_dataset('mnist', train=False, permutation=p) for p in permutations
]
# prepare the model.
mlp = MLP(
DATASET_CONFIGS['mnist']['size']**2,
DATASET_CONFIGS['mnist']['classes'],
hidden_size=args.hidden_size,
hidden_layer_num=args.hidden_layer_num,
hidden_dropout_prob=args.hidden_dropout_prob,
input_dropout_prob=args.input_dropout_prob,
)
# initialize the parameters.
utils.xavier_initialize(mlp)
# prepare the cuda if needed.
if cuda:
mlp.cuda()
# run the experiment.
train(mlp,
train_datasets,
test_datasets,
epochs_per_task=args.epochs_per_task,
batch_size=args.batch_size,
test_size=args.test_size,
consolidate=args.consolidate,
fisher_estimation_sample_size=args.fisher_estimation_sample_size,
lamda=args.lamda,
cuda = torch.cuda.is_available() and args.cuda
task = TASKS[args.task]()
ntm = NTM(
label=task.name,
embedding_size=task.model_input_size,
hidden_size=args.hidden_size,
memory_size=args.memory_size,
memory_feature_size=args.memory_feature_size,
output_size=task.model_output_size,
head_num=args.head_num,
max_shift_size=args.max_shift_size,
)
# initialize the model weights.
utils.xavier_initialize(ntm)
# prepare cuda if needed.
if cuda:
ntm.cuda()
if args.train:
train(
ntm,
task,
model_dir=args.model_dir,
lr=args.lr,
lr_decay=args.lr_decay,
lr_decay_iterations=args.lr_decay_iterations,
weight_decay=args.weight_decay,
batch_size=args.batch_size,
# instantiate the model instance.
wrn = WideResNet(
args.dataset,
dataset_config['size'],
dataset_config['channels'],
dataset_config['classes'],
total_block_number=args.total_block_number,
widen_factor=args.widen_factor,
dropout_prob=args.dropout_prob,
baseline_strides=args.baseline_strides,
baseline_channels=args.baseline_channels,
split_sizes=args.split_sizes,
)
# initialize the weights.
utils.xavier_initialize(wrn)
# prepare cuda if needed.
if cuda:
wrn.cuda()
# run the given command.
if args.test:
utils.load_checkpoint(wrn, args.model_dir, best=True)
utils.validate(wrn,
test_dataset,
test_size=args.test_size,
cuda=cuda,
verbose=True)
else:
train(
"trainB_loader", "validA_loader", "validB_loader"
]
for loader, name in zip(loaders, names):
train_iter = iter(loader)
for _ in range(5):
_, target = train_iter.next()
print(
f'{name}', ': Classes {}, counts: {}'.format(
*np.unique(target.numpy(), return_counts=True)))
##############################
#########Seq Boost############
##############################
model = Net()
utils.xavier_initialize(model)
model = model.to(device)
for name, param in model.named_parameters():
if param.device.type != 'cuda':
print('param {}, not on GPU'.format(name))
optimizer = optim.SGD(model.parameters(),
lr=config['learning_rate'],
momentum=config['momentum'])
mid_model, train_loss, valid_loss = train(model,
trainA_loader,
validA_loader,
batch_size=BATCH_SIZE,
wandb_log=False,
patience=EARLY_STOPPING,
def run(user_args):
# Get full set of user and default args.
args = get_defaults()
args_len = len(args)
args.update(user_args)
assert (len(args) == args_len + 1
) # Ensure there's no typos in the argparser below.
# +1 for dataset arg
print("Running with arguments:")
for arg, value in args.items():
print("{}: {}".format(arg, value))
print()
# Setup CPU or GPU environment
if args['use_cuda'] and not torch.cuda.is_available():
print("Could not access CUDA. Using CPU...")
args['use_cuda'] = False
device = torch.device('cpu')
dtype = torch.FloatTensor
elif args['use_cuda']:
print("Using CUDA.")
device = torch.device('cuda')
dtype = torch.cuda.FloatTensor
else:
print("Using CPU.")
device = torch.device('cpu')
dtype = torch.FloatTensor
print("Initializing model...")
# Define the WGAN
W = WGAN_GP(dtype,
noise_dim=args['noise_dim'],
image_dim=args['image_dim'],
image_channels=args['image_channels'],
disc_channels=args['disc_channels'],
gen_channels=args['gen_channels'],
dataset_name=args['dataset'],
cuda=args['use_cuda'],
tag=args['tag'],
optimizer=args['optim'])
# Initialize weights and parameters
if args['init_method'] == 'gaussian':
utils.gaussian_initialize(W, 0.02)
elif args['init_method'] == 'xavier':
utils.xavier_initialize(W)
print("Initializing dataset...")
# Prepare and load the data
if args['dataset'] == 'mnist':
dataset = dset.MNIST('../data/MNIST',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.ToPILImage(),
transforms.Pad(2),
transforms.ToTensor(),
]))
else:
dataset = GoogleLandmark('../data/train',
args['classes'],
transform=lambda c: transforms.Compose([
transforms.CenterCrop(c),
transforms.Resize(args['image_dim']),
transforms.ToTensor(),
]))
data_loader = DataLoader(dataset,
batch_size=args['batch_size'],
shuffle=True,
drop_last=True,
pin_memory=True if args['use_cuda'] else False)
# Visualize training images in Visdom.
reporter.visualize_images(data_loader.__iter__().next()[0],
'Training data samples',
env=W.name)
# Train!
print("Training model...")
train(W, data_loader, dtype, args)