# Try to load weights if there are any
if LOAD_CHECKPOINTS:
try:
model.load_state_dict(torch.load("checkpoints/params.pt"))
print("Successfully loaded checkpoint")
except:
print("No prior weights found")
# Training loop
model.train()
print("Training...")
for i in range(ITERATIONS):
# Get batch
data, labels = dh.get_batch(BATCH_SIZE)
if CUDA:
data = data.cuda()
labels = labels.cuda()
opt.zero_grad() # reset optimizer
y = model(data) # forward
loss = loss_func(y, labels) # loss calculated
loss.backward()
opt.step()
# Logging, saving, etc.
if LOGGING and i % LOG_INTERVAL == 0:
print("[" + str(i) + "/" + str(ITERATIONS) + "]: " + "Loss: " +
str(loss.item()))
if TESTING and i % TEST_INTERVAL == 0:
