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
