def test_basic_checkpoint(self):
p = torch.tensor([2.0, 1.0, 10.0])
training_steps = 500
model = Net(p)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
tbmodel = torchbearer.Trial(
model,
optim,
loss,
callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')
]).for_train_steps(training_steps).for_val_steps(1)
tbmodel.run(2) # Simulate 2 'epochs'
# Reload
p = torch.tensor([2.0, 1.0, 10.0])
model = Net(p)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
tbmodel = torchbearer.Trial(
model,
optim,
loss,
callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')
]).for_train_steps(training_steps)
tbmodel.load_state_dict(torch.load('test.pt'))
self.assertEqual(len(tbmodel.state[torchbearer.HISTORY]), 2)
self.assertAlmostEqual(model.pars[0].item(), 5.0, places=4)
self.assertAlmostEqual(model.pars[1].item(), 0.0, places=4)
self.assertAlmostEqual(model.pars[2].item(), 1.0, places=4)
import os
os.remove('test.pt')
def main():
fake_parser = FakeArgumentParser(add_help=False, allow_abbrev=False)
add_shared_args(fake_parser)
fake_args, _ = fake_parser.parse_known_args()
parser = argparse.ArgumentParser()
add_shared_args(parser)
add_sub_args(fake_args, parser)
args = parser.parse_args()
trainloader, valloader, testloader = build_dataloaders(args)
args.output.mkdir(exist_ok=True, parents=True)
path = str(args.output)
save_args(args.output)
model = get_model(args.model)()
init_lr, sched = parse_learning_rate_arg(args.learning_rate)
if args.optimiser == 'Adam':
opt = optim.Adam(model.parameters(),
lr=init_lr,
weight_decay=args.weight_decay)
else:
opt = optim.SGD(model.parameters(),
lr=init_lr,
weight_decay=args.weight_decay,
momentum=args.momentum)
callbacks = [
Interval(filepath=path + '/model.{epoch:02d}.pt', period=10),
MostRecent(filepath=path + '/model_final.pt'),
CSVLogger(path + '/train-log.csv'), *sched
]
trial = tb.Trial(model,
opt,
torch.nn.CrossEntropyLoss(),
metrics=['loss', 'acc', 'lr'],
callbacks=callbacks).to(args.device)
trial.with_generators(train_generator=trainloader, val_generator=valloader)
trial.run(epochs=args.epochs, verbose=2)
trial = tb.Trial(model,
criterion=torch.nn.CrossEntropyLoss(),
metrics=['loss', 'acc'],
callbacks=[CSVLogger(path + '/test-log.csv')
]).to(args.device)
trial.with_generators(test_generator=testloader)
trial.predict(verbose=2)
def train(args,
model,
model_loss,
trainloader,
valloader,
epochs,
name='model'):
init_lr, sched = parse_learning_rate_arg(args.learning_rate)
path = str(args.output)
opt = build_optimiser(args, model, init_lr)
callbacks = [
Interval(filepath=path + '/' + name + '.{epoch:02d}.pt', period=10),
MostRecent(filepath=path + '/' + name + '_final.pt'),
CSVLogger(path + '/' + name + '-train-log.csv'), *sched
]
metrics = ['loss', 'lr']
if isinstance(model_loss, nn.CrossEntropyLoss):
metrics.append('acc')
trial = tb.Trial(model,
opt,
model_loss,
metrics=metrics,
callbacks=callbacks).to(args.device)
trial.with_generators(train_generator=trainloader, val_generator=valloader)
trial.run(epochs=epochs, verbose=2)
return trial
def build_trial(args):
model = build_model(args)
loss = build_loss(args)
init_lr, sched = parse_learning_rate_arg(args.learning_rate)
optim = torch.optim.Adam(model.parameters(), lr=init_lr, weight_decay=args.weight_decay)
inv = get_dataset(args.dataset).inv_transform
callbacks = [
Interval(filepath=str(args.output) + '/model_{epoch:02d}.pt', period=args.snapshot_interval),
CSVLogger(str(args.output) + '/log.csv'),
imaging.FromState(tb.Y_PRED, transform=inv).on_val().cache(args.num_reconstructions).make_grid().with_handler(
img_to_file(str(args.output) + '/val_reconstruction_samples_{epoch:02d}.png')),
imaging.FromState(tb.Y_TRUE, transform=inv).on_val().cache(args.num_reconstructions).make_grid().with_handler(
img_to_file(str(args.output) + '/val_samples.png')),
*model.get_callbacks(args),
*sched,
]
if args.variational:
@torchbearer.callbacks.add_to_loss
def add_kld_loss_callback(state):
kl = torch.mean(0.5 * torch.sum(torch.exp(state[LOGVAR]) + state[MU] ** 2 - 1. - state[LOGVAR], 1))
return kl
callbacks.append(add_kld_loss_callback)
trial = tb.Trial(model, optimizer=optim, criterion=loss, metrics=['loss', 'mse', 'lr'], callbacks=callbacks)
trial.with_loader(autoenc_loader)
return trial, model
def train(optimizer_name):
scheduler = torchbearer.callbacks.torch_scheduler.MultiStepLR(milestones=[60, 120, 160], gamma=0.2)
if DATASET == 'CIFAR100':
model = ResNet18(100)
else:
model = ResNet18()
checkpoint = torchbearer.callbacks.ModelCheckpoint(DATASET + "\\" + str(trial_num) + "\\" + optimizer_name + '_checkpoint.pt')
logger = torchbearer.callbacks.CSVLogger(DATASET + "\\" + str(trial_num) + "\\" + optimizer_name + '_log.pt', separator=',', append=True)
if optimizer_name == 'SGD':
optimizer = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9, weight_decay=0.001)
elif optimizer_name == 'Lookahead':
optimizer = Lookahead(torch.optim.SGD(model.parameters(), lr=0.1), la_alpha=0.8, la_steps=5)
elif optimizer_name == 'AdamW':
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1)
elif optimizer_name == 'Polyak':
optimizer = torch.optim.ASGD(model.parameters(), lr=0.3, weight_decay=0.001)
elif optimizer_name == 'Adam':
optimizer = torch.optim.Adam(model.parameters())
elif optimizer_name == 'Lookahead(Adam)':
optimizer = Lookahead(torch.optim.Adam(model.parameters()))
else:
raise ValueError("Optimizer not setup")
loss_function = nn.CrossEntropyLoss()
trial = torchbearer.Trial(model, optimizer, loss_function, metrics=['loss', 'accuracy'], callbacks=[scheduler, checkpoint, logger]).to(device)
trial.with_generators(trainloader, val_generator=valloader)
results.append(trial.run(epochs=NB_EPOCHS))
def on_batch(self, state):
training = state[torchbearer.MODEL].training
state[torchbearer.MODEL].eval()
targets_hot = self._targets_hot(state)
key = self.logit_key
@torchbearer.callbacks.on_sample
def make_eval(_):
state[torchbearer.MODEL].eval()
model = _CAMWrapper(self.input_size,
state[torchbearer.MODEL],
transform=self.in_transform)
trial = torchbearer.Trial(model,
self.optimizer_factory(
filter(lambda p: p.requires_grad,
[model.input_image])),
_cam_loss(key, targets_hot, self.decay),
callbacks=[make_eval])
trial.for_train_steps(self.steps).to(state[torchbearer.DEVICE],
state[torchbearer.DATA_TYPE])
trial.run(verbose=self.verbose)
if training:
state[torchbearer.MODEL].train()
return model.input_image.sigmoid()
def test_only_model(self):
p = torch.tensor([2.0, 1.0, 10.0])
model = Net(p)
tbmodel = torchbearer.Trial(model)
self.assertListEqual(tbmodel.run(), [])
def test_with_loader(self):
p = torch.tensor([2.0, 1.0, 10.0])
training_steps = 2
model = Net(p)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
test_var = {'loaded': False}
def custom_loader(state):
state[torchbearer.X], state[torchbearer.Y_TRUE] = None, None
test_var['loaded'] = True
tbmodel = torchbearer.Trial(
model,
optim,
loss,
callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')
]).for_train_steps(training_steps).for_val_steps(1)
tbmodel.with_loader(custom_loader)
self.assertTrue(not test_var['loaded'])
tbmodel.run(1)
self.assertTrue(test_var['loaded'])
import os
os.remove('test.pt')
def main_wrap(objects, **kwargs):
global acc_plot, data
model = Net(params['objects']) # change to different model here
loss_function = nn.CrossEntropyLoss()
optimiser = optim.Adam(model.parameters(), lr=0.01)
trial = torchbearer.Trial(model,
optimiser,
loss_function,
metrics=['accuracy'],
verbose=0).to(device)
#callbacks=[activation_weights],
dataset = ToyTask(objects, **kwargs) # change to different dataset here
loader = DataLoader(dataset, batch_size=batch_size)
trial.with_generators(train_generator=loader, test_generator=loader)
print("\nTraining:")
# activation_weights.on_end_epoch(None)
trial.run(epochs=train_epochs, verbose=1)
data.append(extract_weights(model))
acc = []
for i in tqdm(range(val_epochs), desc='Testing'):
acc.append(
trial.evaluate(verbose=0,
data_key=torchbearer.TEST_DATA)['test_acc'])
acc_plot.append(np.mean(acc))
def test_no_model(self):
tbmodel = torchbearer.Trial(None)
import warnings
with warnings.catch_warnings(record=True) as w:
tbmodel.run()
self.assertTrue(len(w) == 1)
def test_zero_model(self):
model = Linear(3, 1)
init.constant_(model.weight, 0)
init.constant_(model.bias, 0)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
trial = torchbearer.Trial(model, optim, loss)
trial.with_test_data(torch.rand(10, 3), batch_size=3)
preds = trial.predict()
for i in range(len(preds)):
self.assertAlmostEqual(preds[i], 0)
def test_basic_opt(self):
p = torch.tensor([2.0, 1.0, 10.0])
training_steps = 1000
model = NetWithState(p)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
tbmodel = tb.Trial(model, optim, loss).for_train_steps(training_steps).for_val_steps(1)
tbmodel.run()
self.assertAlmostEqual(model.pars[0].item(), 5.0, places=4)
self.assertAlmostEqual(model.pars[1].item(), 0.0, places=4)
self.assertAlmostEqual(model.pars[2].item(), 1.0, places=4)
def worker():
setup()
print("Rank and node: {}-{}".format(args.rank, platform.node()))
model = ToyModel().to('cpu')
ddp_model = DDP(model)
kwargs = {}
ds = datasets.MNIST('./data/mnist/',
train=True,
download=True,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
train_sampler = torch.utils.data.distributed.DistributedSampler(ds)
train_loader = torch.utils.data.DataLoader(ds,
batch_size=128,
sampler=train_sampler,
**kwargs)
test_ds = datasets.MNIST('./data/mnist',
train=False,
transform=transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307, ), (0.3081, ))
]))
test_sampler = torch.utils.data.distributed.DistributedSampler(test_ds)
test_loader = torch.utils.data.DataLoader(test_ds,
batch_size=128,
sampler=test_sampler,
**kwargs)
loss_fn = nn.CrossEntropyLoss()
optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)
trial = torchbearer.Trial(ddp_model,
optimizer,
loss_fn,
metrics=['loss', 'acc'],
callbacks=[sync, grad, flatten])
trial.with_train_generator(train_loader)
trial.run(10, verbose=2)
print("Model hash: {}".format(hash(model)))
print('First parameter: {}'.format(next(model.parameters())))
cleanup()
def main_wrap(objects, **kwargs):
global acc_plot
model = Net(params['objects'])
loss_function = nn.CrossEntropyLoss()
optimiser = optim.Adam(model.parameters(), lr=0.01)
trial = torchbearer.Trial(model, optimiser, loss_function, metrics=['accuracy']).to(device)
dataset = ToyTask(objects, **kwargs)
loader = DataLoader(dataset, batch_size=batch_size)#, pin_memory=True)
trial.with_generators(train_generator=loader, test_generator=loader)
trial.run(epochs=numepocs, verbose=1)
acc = []
for i in range (200):
acc.append(trial.evaluate(verbose=0, data_key=torchbearer.TEST_DATA)['test_acc'])
acc_plot.append(np.mean(acc))
print(acc_plot)
def test_callbacks(self):
from torch.utils.data import TensorDataset
traingen = TensorDataset(torch.rand(10, 1, 3), torch.rand(10, 1))
valgen = TensorDataset(torch.rand(10, 1, 3), torch.rand(10, 1))
testgen = TensorDataset(torch.rand(10, 1, 3), torch.rand(10, 1))
model = torch.nn.Linear(3, 1)
optim = torch.optim.SGD(model.parameters(), lr=0.01)
cbs = []
cbs.extend([
c.EarlyStopping(),
c.GradientClipping(10, model.parameters()),
c.Best('test.pt'),
c.MostRecent('test.pt'),
c.ReduceLROnPlateau(),
c.CosineAnnealingLR(0.1, 0.01),
c.ExponentialLR(1),
c.Interval('test.pt'),
c.CSVLogger('test_csv.pt'),
c.L1WeightDecay(),
c.L2WeightDecay(),
c.TerminateOnNaN(monitor='fail_metric')
])
trial = torchbearer.Trial(model,
optim,
torch.nn.MSELoss(),
metrics=['loss'],
callbacks=cbs)
trial = trial.with_generators(traingen, valgen, testgen)
trial.run(2)
trial.predict()
trial.evaluate(data_key=torchbearer.TEST_DATA)
trial.evaluate()
import os
os.remove('test.pt')
os.remove('test_csv.pt')
def on_batch(self, state):
training = state[torchbearer.MODEL].training
@torchbearer.callbacks.on_sample
def make_eval(_):
state[torchbearer.MODEL].eval()
@torchbearer.callbacks.add_to_loss
def loss(state):
return -self.criterion(state)
model = _Wrapper(self.image, state[torchbearer.MODEL])
trial = torchbearer.Trial(model,
self.optimizer,
callbacks=[make_eval, loss])
trial.for_train_steps(self.steps).to(state[torchbearer.DEVICE],
state[torchbearer.DATA_TYPE])
trial.run(verbose=self.verbose)
if training:
state[torchbearer.MODEL].train()
return model.image.get_valid_image()
def on_batch(self, state):
training = state[torchbearer.MODEL].training
state[torchbearer.MODEL].eval()
targets_hot = self._targets_hot(state)
to_prob = (lambda x: x.exp()) if self.prob_key is None else (
lambda x: x)
key = self.logit_key if self.prob_key is None else self.prob_key
model = _CAMWrapper(self.input_size, state[torchbearer.MODEL])
trial = torchbearer.Trial(
model,
self.optimizer_factory(
filter(lambda p: p.requires_grad, [model.input_batch])),
_cam_loss(key, to_prob, targets_hot, self.decay))
trial.for_train_steps(self.steps).to(state[torchbearer.DEVICE],
state[torchbearer.DATA_TYPE])
trial.run(verbose=self.verbose)
if training:
state[torchbearer.MODEL].train()
return model.input_batch.squeeze(0)
def build_test_trial(args):
model = build_model(args).to(args.device)
inv = get_dataset(args.dataset).inv_transform
callbacks = [
CSVLogger(str(args.output) + '/log.csv'),
imaging.FromState(tb.Y_PRED, transform=inv).on_test().cache(args.num_reconstructions).make_grid().with_handler(
img_to_file(str(args.output) + '/test_reconstruction_samples.png')),
imaging.FromState(tb.Y_TRUE, transform=inv).on_test().cache(args.num_reconstructions).make_grid().with_handler(
img_to_file(str(args.output) + '/test_samples.png'))
]
metrics = ['mse'] #Additional metrics: ChamferMetric(), ModifiedHausdorffMetric()
if args.classifier_weights:
classifier = get_classifier_model(args.classifier_model)().to(args.device)
state = torch.load(args.classifier_weights, map_location=args.device)
classifier.load_state_dict(state[tb.MODEL])
metrics.append(ClassificationMetric(classifier))
trial = tb.Trial(model, metrics=metrics, callbacks=callbacks)
trial.with_loader(autoenc_loader)
return trial, model
# gaussian blur the image with the chosen sigma value
blurredImage = blur(testset.test_data[i].numpy(), sigma_values[j])
# update the training set with the blurred image
testset.test_data[i] = torch.from_numpy(blurredImage)
trainloader = DataLoader(trainset, batch_size=128, shuffle=True)
testloader = DataLoader(testset, batch_size=128, shuffle=True)
#test the controlled network
model = torch.load('save_ann.pkl')
loss_function = nn.CrossEntropyLoss()
# live_loss_plot = LiveLossPlot()
optimiser = optim.SGD(model.parameters(), lr=0.1, momentum=0.5)
# trial = torchbearer.Trial(model, optimiser, loss_function, callbacks=[live_loss_plot], metrics=['loss', 'accuracy']).to(device)
trial = torchbearer.Trial(model,
optimiser,
loss_function,
metrics=['loss', 'accuracy']).to(device)
trial.with_generators(trainloader, test_generator=testloader)
results = trial.evaluate(data_key=torchbearer.TEST_DATA)
print(results)
blur_acc_control.append(100.0 * results["test_acc"])
#test the sleeping network
model = torch.load('save_sleep.pkl')
loss_function = nn.CrossEntropyLoss()
# live_loss_plot = LiveLossPlot()
optimiser = optim.SGD(model.parameters(), lr=0.1, momentum=0.5)
# trial = torchbearer.Trial(model, optimiser, loss_function, callbacks=[live_loss_plot], metrics=['loss', 'accuracy']).to(device)
trial = torchbearer.Trial(model,
optimiser,
loss_function,
return state[tb.MODEL].x.data
@tb.callbacks.on_step_training
def greedy_update(state):
if state[tb.MODEL].x > 1:
state[tb.MODEL].x.data.fill_(1)
elif state[tb.MODEL].x < -1:
state[tb.MODEL].x.data.fill_(-1)
training_steps = 6000000
model = Online()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99])
tbtrial = tb.Trial(model, optim, loss, [est()], pass_state=True, callbacks=[greedy_update, TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)])
tbtrial.for_train_steps(training_steps).run()
model = Online()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True)
tbtrial = tb.Trial(model, optim, loss, [est()], pass_state=True, callbacks=[greedy_update, TensorBoard(comment='amsgrad', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)])
tbtrial.for_train_steps(training_steps).run()
model = Stochastic()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99])
tbtrial = tb.Trial(model, optim, loss, [est()], callbacks=[greedy_update, TensorBoard(comment='adam', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)])
tbtrial.for_train_steps(training_steps).run()
model = Stochastic()
optim = torch.optim.Adam(model.parameters(), lr=0.001, betas=[0.9, 0.99], amsgrad=True)
tbtrial = tb.Trial(model, optim, loss, [est()], callbacks=[greedy_update, TensorBoard(comment='amsgrad', write_graph=False, write_batch_metrics=True, write_epoch_metrics=False)])
def test_no_model(self):
trial = torchbearer.Trial(None)
trial.run()
self.assertTrue(torchbearer.trial.MockModel()(torch.rand(1)) is None)
def main():
fake_parser = FakeArgumentParser(add_help=False, allow_abbrev=False)
add_shared_args(fake_parser)
fake_args, _ = fake_parser.parse_known_args()
parser = argparse.ArgumentParser()
add_shared_args(parser)
add_sub_args(fake_args, parser)
args = parser.parse_args()
rndtf = transforms.Compose([
transforms.RandomAffine(10.0,
translate=(0.1, 0.1),
scale=(0.95, 1.01),
shear=1,
fillcolor=None),
transforms.Lambda(lambda x: x * (1 + (torch.rand_like(x) - 0.5) / 10))
])
args.additional_transforms = transforms.Lambda(lambda x:
(rndtf(x), rndtf(x)))
orig_batch_size = args.batch_size
args.batch_size = args.barlow_batch_size
trainloader, valloader, testloader = build_dataloaders(args)
args.batch_size = orig_batch_size
args.output.mkdir(exist_ok=True, parents=True)
save_args(args.output)
model = get_model(args.model)()
btmodel = BarlowTwinsModel(model)
model_loss = nn.CrossEntropyLoss()
btmodel_loss = BarlowTwinsLoss()
train(args,
btmodel,
btmodel_loss,
trainloader,
valloader,
args.barlow_epochs,
name='btmodel')
model.lock_features(not args.finetune)
trainloader, valloader, testloader = build_dataloaders(
args) # reload data with other batch size
train(args,
model,
model_loss,
trainloader,
valloader,
args.epochs,
name='model')
trial = tb.Trial(model,
criterion=torch.nn.CrossEntropyLoss(),
metrics=['loss', 'acc'],
callbacks=[CSVLogger(str(args.output) + '/test-log.csv')
]).to(args.device)
trial.with_generators(test_generator=testloader)
trial.predict(verbose=2)
def test_no_train_steps(self):
trial = torchbearer.Trial(None)
trial.for_val_steps(10)
trial.run()
# Model and optimizer
model = GAN()
optim = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.5, 0.999))
@tb.metrics.running_mean
@tb.metrics.mean
class g_loss(tb.metrics.Metric):
def __init__(self):
super().__init__('g_loss')
def process(self, state):
return state[G_LOSS]
@tb.metrics.running_mean
@tb.metrics.mean
class d_loss(tb.metrics.Metric):
def __init__(self):
super().__init__('d_loss')
def process(self, state):
return state[D_LOSS]
torchbearertrial = tb.Trial(model, optim, criterion=None, metrics=['loss', g_loss(), d_loss()],
callbacks=[loss_callback, saver_callback], pass_state=True)
torchbearertrial.with_train_generator(dataloader)
torchbearertrial.to(device)
torchbearertrial.run(epochs=200)
tb.LOSS, DISC_OPT)
def closure(state):
closure_gen(state)
state[GEN_OPT].step()
closure_disc(state)
state[DISC_OPT].step()
from torchbearer.metrics import mean, running_mean
metrics = ['loss', mean(running_mean(D_LOSS)), mean(running_mean(G_LOSS))]
trial = tb.Trial(generator,
None,
criterion=gen_crit,
metrics=metrics,
callbacks=[saver_callback])
trial.with_train_generator(dataloader, steps=200000)
trial.to(device)
new_keys = {
DISC_MODEL: discriminator.to(device),
DISC_OPT: optimizer_D,
GEN_OPT: optimizer_G,
DISC_CRIT: disc_crit
}
trial.state.update(new_keys)
trial.with_closure(closure)
trial.run(epochs=1)
x = [5,0,1]
"""
out = torch.zeros_like(self.pars)
out[0] = self.pars[0] - 5
out[1] = self.pars[1]
out[2] = self.pars[2] - 1
return torch.sum(out**2)
def forward(self, _, state):
state[ESTIMATE] = self.pars.detach().unsqueeze(1)
return self.f()
def loss(y_pred, y_true):
return y_pred
p = torch.tensor([2.0, 1.0, 10.0])
training_steps = 50000
model = Net(p)
optim = torch.optim.SGD(model.parameters(), lr=0.0001)
tbtrial = tb.Trial(model,
optim,
loss, [tb.metrics.running_mean(ESTIMATE, dim=1), 'loss'],
pass_state=True)
tbtrial.for_train_steps(training_steps).to('cuda')
tbtrial.run()
print(list(model.parameters())[0].data)
# utilize the GPU if one exists, otherwise use the CPU
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# initialize control network
control_model = NetworkControl(784, 1200, 10)
# define the loss function
loss_function = nn.CrossEntropyLoss()
# define the optimiser, learning rate and momentum
optimiser = optim.SGD(model.parameters(), lr=0.1, momentum=0.5)
# create a live loss plot of accuracy and loss after each training epoch
plot = LiveLossPlot()
# train the network using model, optimiser and loss function for 2 epochs
trial = torchbearer.Trial(control_model, optimiser, loss_function, callbacks=[plot], metrics=['loss', 'accuracy']).to(device)
trial.with_generators(trainloader, test_generator=testloader)
trial.run(epochs=2)
# save the trained model
torch.save(model.state_dict(),"control.pth")
# print the accuracy and loss of control network on the MNIST test data
results = trial.evaluate(data_key=torchbearer.TEST_DATA)
print("Test Loss:\t", results.get("test_loss"))
print("Test Accuracy:\t", results.get("test_acc"))
from scipy.ndimage import gaussian_filter
import numpy as np
from matplotlib import pyplot as plt
import random
