def main():
device = torch.device('cpu')
transform=transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))]) # MNIST-specific values
train_set = datasets.MNIST('mnist_data', train=True, transform=transform)
test_set = datasets.MNIST('mnist_data', train=False, transform=transform)
train_loader = torch.utils.data.DataLoader(train_set, batch_size=BATCH_SIZE)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=TEST_BATCH_SIZE)
model = Model()
model = model.to(device)
optimizer = torch.optim.Adadelta(model.parameters(), lr=LEARNING_RATE)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, gamma=GAMMA, step_size=1)
for i in range(EPOCHS):
train(model, train_loader, optimizer, device)
scheduler.step()
test(model, test_loader, device)
torch.save(model.state_dict(), "parameters.pt")
raw3_percent.append(percent_AEE)
print(f"Raw3/AEE = {np.mean(raw3_AEE)}")
print(f"Raw3/percent = {np.mean(raw3_percent)}")
writer.add_scalar('Raw3/AEE', np.mean(raw3_AEE), epoch)
writer.add_scalar('Raw3/percent', np.mean(raw3_percent), epoch)
print(
f"Test/AEE = {np.mean([np.mean(raw1_AEE), np.mean(raw2_AEE), np.mean(raw3_AEE)])}"
)
print(
f"Test/percent = {np.mean([np.mean(raw1_percent), np.mean(raw2_percent), np.mean(raw3_percent)])}"
)
writer.add_scalar(
'Test/AEE',
np.mean([np.mean(raw1_AEE),
np.mean(raw2_AEE),
np.mean(raw3_AEE)]), epoch)
writer.add_scalar(
'Test/percent',
np.mean([
np.mean(raw1_percent),
np.mean(raw2_percent),
np.mean(raw3_percent)
]), epoch)
model.train()
writer.close()
torch.save(model.state_dict(), models_path / "only_kitty_with_rotation.pth")
