N, D_in, H, D_out = 64, 1000, h, 10 # Create random Tensors to hold inputs and outputs x = torch.randn(N, D_in) y = torch.randn(N, D_out) model = TwoLayerNet(D_in, H, D_out) x, y, model = x.to(device), y.to(device), model.to(device) layers = [model.linear1, model.linear2] stats = CheckLayerSat('regression/h{}'.format(h), layers) loss_fn = torch.nn.MSELoss(size_average=False) optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9) steps_iter = trange(2000, desc='steps', leave=True, position=0) steps_iter.write("{:^80}".format( "Regression - TwoLayerNet - Hidden layer size {}".format(h))) for _ in steps_iter: y_pred = model(x) loss = loss_fn(y_pred, y) steps_iter.set_description('loss=%g' % loss.data) optimizer.zero_grad() loss.backward() optimizer.step() stats.saturation() steps_iter.write('\n') stats.close() steps_iter.close()
def train(network, dataset, test_set, logging_dir, batch_size): network.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(network.parameters()) #stats = CheckLayerSat(logging_dir, network, log_interval=len(dataset)//batch_size) stats = CheckLayerSat(logging_dir, network, log_interval=60, sat_method='cumvar99', conv_method='mean') epoch_acc = 0 thresh = 0.95 epoch = 0 total = 0 correct = 0 value_dict = None while epoch <= 20: print('Start Training Epoch', epoch, '\n') start = t.time() epoch_acc = 0 train_loss = 0 total = 0 correct = 0 network.train() for i, data in enumerate(dataset): step = epoch * len(dataset) + i inputs, labels = data inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = network(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() train_loss += loss.item() _, predicted = outputs.max(1) total += labels.size(0) correct += predicted.eq(labels).sum().item() #if i % 2000 == 1999: # print every 2000 mini-batches print(i, 'of', len(dataset), 'acc:', correct / total) # display layer saturation levels end = t.time() stats.saturation() test_loss, test_acc = test(network, test_set, criterion, stats, epoch) epoch_acc = correct / total print('Epoch', epoch, 'finished', 'Acc:', epoch_acc, 'Loss:', train_loss / total, '\n') stats.add_scalar('train_loss', train_loss / total, epoch) # optional stats.add_scalar('train_acc', epoch_acc, epoch) # optional value_dict = record_metrics(value_dict, stats.logs, epoch_acc, train_loss / total, test_acc, test_loss, epoch, (end - start) / total) log_to_csv(value_dict, logging_dir) epoch += 1 stats.close() # test_stats.close() return criterion