def train(epochs, train_loader, dev_loader, lr, seed, log_interval,
output_dir):
"""Train the model. Store snapshot models in the output_dir alongside
evaluations on the dev set after each epoch
"""
model = Net()
optimizer = optim.Adam(model.parameters(), lr=lr)
measure_size(model)
use_cuda = torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu")
print("Using device: ", device)
if use_cuda:
torch.cuda.manual_seed(seed)
else:
torch.manual_seed(seed)
#torch.backends.cudnn.benchmark = False
#torch.backends.cudnn.deterministic = True
model.to(device)
for epoch in range(1, epochs):
model.train()
total_loss = 0.0
for batch_idx, (data, target) in enumerate(train_loader):
if use_cuda:
data, target = data.to(device), target.to(device)
data = data.unsqueeze_(1)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
total_loss += loss.item()
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
print("Total loss = %.6f" % (total_loss / len(train_loader.dataset)))
test(model, dev_loader,
os.path.join(output_dir, 'dev-eer-' + str(epoch)))
torch.save(model, os.path.join(output_dir,
'iter' + str(epoch) + '.mdl'))
class Trainer(object):
def __init__(self, args):
self.args = args
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.prepare_data()
self.setup_train()
def prepare_data(self):
train_val = MnistDataset(
self.args.train_image_file,
self.args.train_label_file,
transform=transforms.Compose([ToTensor()]),
)
train_len = int(0.8 * len(train_val))
train_ds, val_ds = torch.utils.data.random_split(
train_val, [train_len, len(train_val) - train_len]
)
print("Train {}, val {}".format(len(train_ds), len(val_ds)))
self.train_loader = torch.utils.data.DataLoader(
train_ds,
batch_size=self.args.batch_size,
collate_fn=collate_fn,
shuffle=True,
)
self.val_loader = torch.utils.data.DataLoader(
val_ds,
batch_size=self.args.batch_size,
collate_fn=collate_fn,
shuffle=False,
)
def setup_train(self):
self.model = Net().to(self.device)
self.optimizer = torch.optim.SGD(self.model.parameters(), lr=self.args.lr)
self.criterion = nn.CrossEntropyLoss().to(self.device)
if not os.path.isdir(self.args.ckpt):
os.mkdir(self.args.ckpt)
def train_one_epoch(self):
train_loss = 0.0
self.model.train()
for i, sample in enumerate(self.train_loader):
X, Y_true = sample["X"].to(self.device), sample["Y"].to(self.device)
self.optimizer.zero_grad()
output = self.model(X)
loss = self.criterion(output, Y_true)
loss.backward()
self.optimizer.step()
train_loss += loss.item()
return train_loss / len(self.train_loader)
def evaluate(self):
val_loss = 0.0
self.model.eval()
predicts = []
truths = []
with torch.no_grad():
for i, sample in enumerate(self.val_loader):
X, Y_true = sample["X"].to(self.device), sample["Y"].to(self.device)
output = self.model(X)
loss = self.criterion(output, Y_true)
val_loss += loss.item()
predicts.append(torch.argmax(output, dim=1))
truths.append(Y_true)
predicts = torch.cat(predicts, dim=0)
truths = torch.cat(truths, dim=0)
acc = torch.sum(torch.eq(predicts, truths))
return acc / len(predicts), val_loss / (len(self.val_loader))
def run(self):
min_loss = 10e4
max_acc = 0
for epoch in range(self.args.epochs):
train_loss = self.train_one_epoch()
val_acc, val_loss = self.evaluate()
if val_acc > max_acc:
max_acc = val_acc
torch.save(
self.model.state_dict(),
os.path.join(
self.args.ckpt,
"{}_{}_{:.4f}.pth".format(self.args.name, epoch, max_acc),
),
)
print(
"Epoch {}, loss {:.4f}, val_acc {:.4f}".format(
epoch, train_loss, val_acc
)
)