def train(cnn, epochs=80, learn_rate=0.001, batch_size=100, gpu=True):
"""
Train a regression CNN. Note that you do not need this function.
Included for refrence.
"""
if gpu:
cnn.cuda()
# Set up L2 loss
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(cnn.parameters(), lr=learn_rate)
# Loading & transforming data
(x_train, y_train), (x_test, y_test) = load_cifar10()
train_rgb, train_grey = process(x_train, y_train)
test_rgb, test_grey = process(x_test, y_test)
print(train_rgb)
print(len(train_rgb[0]))
print(len(train_rgb[0][0]))
print(len(train_rgb[0][0][0]))
input()
print("Beginning training ...")
for epoch in range(epochs):
# Train the Model
cnn.train() # Change model to 'train' mode
for i, (xs, ys) in enumerate(get_batch(train_grey,
train_rgb,
batch_size)):
images, labels = get_torch_vars(xs, ys, gpu)
# Forward + Backward + Optimize
optimizer.zero_grad()
outputs = cnn(images)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
print('Epoch [%d/%d], Loss: %.4f' % (epoch+1, epochs, loss.data[0]))
# Evaluate the model
cnn.eval() # Change model to 'eval' mode (BN uses moving mean/var).
losses = []
for i, (xs, ys) in enumerate(get_batch(test_grey,
test_rgb,
batch_size)):
images, labels = get_torch_vars(xs, ys, gpu)
outputs = cnn(images)
val_loss = criterion(outputs, labels)
losses.append(val_loss.data[0])
val_loss = np.mean(losses)
print('Epoch [%d/%d], Val Loss: %.4f' % (epoch+1, epochs, val_loss))
# Save the Trained Model
torch.save(cnn.state_dict(), 'regression_cnn_k%d_f%d.pkl' % (
args.kernel, args.num_filters))
args = parser.parse_args()
npr.seed(0)
cnn = RegressionCNN(args.kernel, args.num_filters)
# Uncomment to train. You do not need this for the assignment.
# Included for completeness
#train(cnn); exit(0)
print("Loading weights...")
checkpoint = torch.load('weights/regression_cnn_k%d_f%d.pkl' %
(args.kernel, args.num_filters),
map_location=lambda storage, loc: storage)
cnn.load_state_dict(checkpoint)
print("Loading data...")
(x_train, y_train), (x_test, y_test) = load_cifar10()
test_rgb, test_grey = process(x_test, y_test)
# Create output folder if not created
if not os.path.exists("outputs"):
os.makedirs("outputs")
print("Generating predictions...")
grey = test_grey[:15]
gtrgb = test_rgb[:15]
predrgb = cnn(Variable(torch.from_numpy(grey).float()))
predrgb = predrgb.data.numpy()
plot(grey, gtrgb, predrgb, "outputs/regression_output.png")