def normal_train(args, loader_train, loader_test, dtype): model = ConvNet() model = model.type(dtype) model.train() loss_f = nn.CrossEntropyLoss() SCHEDULE_EPOCHS = [15] learning_rate = 0.01 for num_epochs in SCHEDULE_EPOCHS: print('\nTraining %d epochs with learning rate %.4f' % (num_epochs, learning_rate)) optimizer = optim.Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): print('\nTraining epoch %d / %d ...\n' % (epoch + 1, num_epochs)) # print(model.training) for i, (X_, y_) in enumerate(loader_train): X = Variable(X_.type(dtype), requires_grad=False) y = Variable(y_.type(dtype), requires_grad=False).long() preds = model(X) loss = loss_f(preds, y) if (i + 1) % args.print_every == 0: print('Batch %d done, loss = %.7f' % (i + 1, loss.item())) optimizer.zero_grad() loss.backward() optimizer.step() print('Batch %d done, loss = %.7f' % (i + 1, loss.item())) test(model, loader_test, dtype) learning_rate *= 0.1 return model
def unrolled(args, loader_train, loader_test, dtype): model = ConvNet() model = model.type(dtype) model.train() SCHEDULE_EPOCHS = [50, 50] learning_rate = 5e-4 for num_epochs in SCHEDULE_EPOCHS: print('\nTraining %d epochs with learning rate %.7f' % (num_epochs, learning_rate)) optimizer = optim.Adam(model.parameters(), lr=learning_rate) for epoch in range(num_epochs): print('\nTraining epoch %d / %d ...\n' % (epoch + 1, num_epochs)) # print(model.training) for i, (X_, y_) in enumerate(loader_train): X = Variable(X_.type(dtype), requires_grad=False) y = Variable(y_.type(dtype), requires_grad=False) loss = cw_train_unrolled(model, X, y, dtype) optimizer.zero_grad() loss.backward() optimizer.step() if (i + 1) % args.print_every == 0: print('Batch %d done, loss = %.7f' % (i + 1, loss.item())) test(model, loader_test, dtype) print('Batch %d done, loss = %.7f' % (i + 1, loss.item())) learning_rate *= 0.1 return model