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")