def train(model, train_data, val_data, params):
optimizer = torch.optim.Adam(model.parameters(),
lr=params['lr'],
weight_decay=params['decay'])
# optimizer = torch.optim.SGD(model.parameters(), lr=params['lr'], momentum=params['p_i'], dampening=params['p_i'])
scheduler = lr_scheduler.StepLR(optimizer,
step_size=params['step_size'],
gamma=params['gamma'])
pruner = Pruner(train_data, optimizer, params)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=params['batch_size'],
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=params['batch_size'],
shuffle=True)
# output_file = 'bs_'+str(batch_size)+'_lr_'+str(learning_rate)+'_wd_'+str(decay)+'.csv'
# with open(output_file, 'w', newline="") as f_out:
# writer = csv.writer(f_out, delimiter=',')
# writer.writerow(["Epoch",'Train Loss','Val Loss','Val Acc'])
df = pd.DataFrame(index=list(range(params['epochs'])),
columns=["trainacc", "trainloss", "valacc", "valloss"])
# fig = plt.gcf()
# fig.show()
# fig.canvas.draw()
for epoch in range(
params['epochs']): # loop over the dataset multiple times
t1 = time.perf_counter()
model.train()
train_losses, train_accs = [], []
acc = 0
for batch, (x_train, y_train) in enumerate(train_loader):
model.zero_grad()
pred, std = model(x_train)
if epoch > params['prune_milestones'][0]:
loss = F.cross_entropy(pred, y_train)
else:
d, yd = model.sample_sigma(pred, std, y_train)
loss = F.cross_entropy(pred, y_train) + loss_l2(d, yd)
loss.backward()
optimizer.step()
# optimizer.param_groups[0]['lr'] = get_lr(epoch, params['lr'], params['gamma'])
# optimizer.param_groups[0]['momentum'] = get_momentum(epoch, params['p_i'], params['p_f'], params['T'])
# optimizer.param_groups[0]['dampening'] = get_momentum(epoch, params['p_i'], params['p_f'], params['T'])
acc = (pred.argmax(dim=-1) == y_train).to(torch.float32).mean()
train_accs.append(acc.mean().item())
train_losses.append(loss.item())
with torch.no_grad():
model.eval()
val_losses, val_accs = [], []
acc = 0
for i, (x_val, y_val) in enumerate(val_loader):
val_pred, val_std = model(x_val)
if epoch > params['prune_milestones'][0]:
val_loss = F.cross_entropy(val_pred, y_val)
else:
val_d, val_yd = model.sample_sigma(val_pred, val_std,
y_val)
val_loss = F.cross_entropy(val_pred, y_val) + loss_l2(
val_d, val_yd)
# val_d, val_yd = model.sample_sigma(val_pred, val_std,y_val)
# loss = F.cross_entropy(val_pred,y_val) + loss_l2(val_d,val_yd)
acc = (val_pred.argmax(dim=-1) == y_val).to(
torch.float32).mean()
val_losses.append(val_loss.item())
val_accs.append(acc.mean().item())
train_loader = pruner.prune_dataloaders(model, train_loader)
scheduler.step()
print(
'Decay: {}, Epoch: {}, Loss: {}, VAccuracy: {}, Ignore_list_Size: {}, LR: {}'
.format(params['decay'], epoch, np.mean(val_losses),
np.mean(val_accs), len(pruner.ignore_list),
optimizer.param_groups[0]['lr']))
df.loc[epoch, 'trainacc'] = np.mean(train_accs)
df.loc[epoch, 'trainloss'] = np.mean(train_losses)
df.loc[epoch, 'valacc'] = np.mean(val_accs)
df.loc[epoch, 'valloss'] = np.mean(val_losses)
# plt.plot(df[:epoch+1]['valloss'])
# plt.plot(df[:epoch+1]['trainloss'])
# plt.pause(0.01) # I ain't needed!!!
# fig.canvas.draw()
fig1 = plt.figure()
ax1 = fig1.add_subplot()
ax1.plot(df['valloss'])
ax1.plot(df['trainloss'])
ax1.grid()
ax1.set_xlabel('Epoch')
fig1.show()
return model, df
def train_prune(model, train_data, val_data, params):
init_model_dict = model.state_dict()
optimizer = torch.optim.Adam(model.parameters(),
lr=params['lr'],
weight_decay=params['decay'])
scheduler = lr_scheduler.StepLR(optimizer,
step_size=params['step_size'],
gamma=params['gamma'])
pruner = Pruner(train_data, optimizer, params)
train_loader = torch.utils.data.DataLoader(train_data,
batch_size=params['batch_size'],
shuffle=True)
val_loader = torch.utils.data.DataLoader(val_data,
batch_size=params['batch_size'],
shuffle=True)
# output_file = 'bs_'+str(batch_size)+'_lr_'+str(learning_rate)+'_wd_'+str(decay)+'.csv'
# with open(output_file, 'w', newline="") as f_out:
# writer = csv.writer(f_out, delimiter=',')
# writer.writerow(["Epoch",'Train Loss','Val Loss','Val Acc'])
df = pd.DataFrame(index=list(range(params['epochs'])),
columns=["trainacc", "trainloss", "valacc", "valloss"])
fig = plt.gcf()
fig.show()
fig.canvas.draw()
epoch_trainaccs, epoch_valaccs = [], []
epoch_trainloss, epoch_valloss = [], []
for epoch in range(
params['epochs']): # loop over the dataset multiple times
t1 = time.perf_counter()
model.train()
train_losses, train_accs = [], []
acc = 0
for batch, (x_train, y_train) in enumerate(train_loader):
model.zero_grad()
pred, std = model(x_train)
if (epoch > params['prune_milestones'][0]):
d, yd = model.sample_sigma(pred, std, y_train)
loss = F.cross_entropy(pred, y_train) + loss_l2(d, yd)
else:
loss = F.cross_entropy(pred, y_train)
loss.backward()
optimizer.step()
acc = (pred.argmax(dim=-1) == y_train).to(torch.float32).mean()
train_accs.append(acc.mean().item())
train_losses.append(loss.item())
with torch.no_grad():
model.eval()
val_losses, val_accs = [], []
acc = 0
for i, (x_val, y_val) in enumerate(val_loader):
val_pred, val_std = model(x_val)
val_d, val_yd = model.sample_sigma(val_pred, val_std, y_val)
loss = F.cross_entropy(val_pred, y_val) + loss_l2(
val_d, val_yd)
acc = (val_pred.argmax(dim=-1) == y_val).to(
torch.float32).mean()
val_losses.append(loss.item())
val_accs.append(acc.mean().item())
train_loader = pruner.prune_dataloaders(model, train_loader)
scheduler.step()
print(
'Decay: {}, Epoch: {}, Loss: {}, TAccuracy: {}, Ignore_list_Size: {}, LR: {}'
.format(params['decay'], epoch, np.mean(val_losses),
np.mean(train_accs), len(pruner.ignore_list),
optimizer.param_groups[0]['lr']))
df.loc[epoch, 'trainacc'] = np.mean(train_accs)
df.loc[epoch, 'trainloss'] = np.mean(train_losses)
df.loc[epoch, 'valacc'] = np.mean(val_accs)
df.loc[epoch, 'valloss'] = np.mean(val_losses)
plt.plot(df[:epoch + 1]['valloss'])
plt.plot(df[:epoch + 1]['trainloss'])
plt.pause(0.01) # I ain't needed!!!
fig.canvas.draw()
epoch_trainaccs.append(np.mean(train_accs))
epoch_valaccs.append(np.mean(val_accs))
epoch_trainloss.append(np.mean(train_losses))
epoch_valloss.append(np.mean(val_losses))
_results = [
time.perf_counter() - t1, epoch,
np.mean(train_losses),
np.mean(val_losses),
np.mean(val_accs)
]
# with open(output_file, 'a', newline="") as f_out:
# writer = csv.writer(f_out, delimiter=',')
# writer.writerow(_results)
# with torch.no_grad():
# unc = []
# model.eval()
# for b,(x,y) in enumerate(train_loader):
# pred,s = model(x)
# c = pred.argmax(dim=1)
# for i in range(len(c)):
# unc.append(torch.sqrt(F.softplus(s[i,c[i]])).item())
# fig = plt.figure()
# ax = fig.add_subplot()
# ax.hist(unc,100)
# ax.set_title('Aleatoric Uncertainty distribution of training data after pruning')
return model, df, train_loader