예제 #1
0
파일: training.py 프로젝트: stas97/bsc_lstm
def one_epoch_training(model, loss, optimiser, scheduler, device, train_gen, test_gen,
                       timesteps, batch_size):
    t0 = time.time()
    scheduler = scheduler
    optimiser = optimiser
    loss_fn = loss

    # Track losses:
    loss_vals_train = np.zeros((len(train_gen)))
    loss_vals_test = np.zeros((len(test_gen)))

    total_time = 0
    batch_nr = 0
    for batch, labels in train_gen:
        model.batch_size = batch_size
        batch = batch.view(timesteps, batch_size, -1)
        # print(batch.shape)
        labels = labels.float()
        # Transfer to GPU
        batch, labels = batch.to(device), labels.to(device)
        optimiser.zero_grad()
        preds = model(batch)
        loss = loss_fn(preds, labels)
        loss_vals_train[batch_nr] = loss.item()
        loss.backward()
        optimiser.step()

        if scheduler is not None:
            scheduler.step()
        batch_nr += 1

    batch_nr = 0
    for batch, labels in test_gen:
        model.batch_size = 1
        batch = batch.view(timesteps, 1, -1)
        labels = labels.float()
        # Transfer to GPU
        batch, labels = batch.to(device), labels.to(device)
        optimiser.zero_grad()
        y_pred_test = model(batch)
        loss = loss_fn(y_pred_test, labels)
        loss_vals_test[batch_nr] = loss.item()
        batch_nr += 1
    error = np.mean(loss_vals_test)
    print('The epoch took {} seconds'.format(time.time() - t0))
    total_time += time.time() - t0

    return error, model
예제 #2
0
    def _predict_n_forward(self, dataset, model_key, window_normalisation=True):
        """
        Idea:
            Do a dataset with the longest tiemesteps
            feed it to run_predictions
            inside run predictions split it up to accommodate shorter timestep sequences
         """
        timesteps = self.model_config[model_key]["timesteps"]
        model = self.models_dict[model_key]
        n_forward = self.model_config[model_key]['num_forward']
        dt = dataset.get_test_data(timesteps, window_normalisation, n_forward)[0]
        predicted = []
        predicted_denormalized = []
        skip = self.most_timesteps - timesteps

        for i in range(int(len(dt) / self.prediction_length)):
            if skip + i * self.prediction_length < dt.shape[0]:
                curr_frame = dt[skip + i * self.prediction_length]
                curr_frame = torch.from_numpy(curr_frame).type(torch.Tensor).detach().view(timesteps, 1, -1)
                prediction = model(curr_frame)[0, 0].detach()
                predicted.append(prediction.cpu().numpy())
                predicted_denormalized.append(
                    denormalise("window", dataset.w_normalisation_p0_test[i * self.prediction_length][0],
                                prediction.cpu().numpy(), None)
                )
        return predicted, predicted_denormalized
예제 #3
0
def test(model, dataset):
    avg_loss = tf.keras.metrics.Mean('loss', dtype=tf.float32)
    avg_accuracy = tf.keras.metrics.Mean('accuracy', dtype=tf.float32)

    for images, labels in dataset:
        logits = model(images, training=False)

        avg_loss(loss_fn(labels, logits))
        avg_accuracy(accuracy_fn(labels, logits))

    print('# test_loss :', avg_loss.result().numpy(), ', test accurary :', avg_accuracy.result().numpy())
예제 #4
0
    def train_step(images, labels):
        with tf.GradientTape() as tape:
            logits = model(images, training=True)

            loss = loss_fn(labels, logits)
            accuracy = accuracy_fn(labels, logits)

        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))

        return loss, accuracy
예제 #5
0
def evaluate(model, val_loader, criterion):
    model.eval()
    epoch_loss = 0
    epoch_acc = 0
    with torch.no_grad():
        for X, Y in val_loader:
            X = X.to(c.device)
            Y = Y.to(c.device)
            outputs = model(X)
            epoch_acc += torch.sum((torch.round(outputs) == Y)).float() / Y.size(0)
            loss = criterion(outputs, Y)
            epoch_loss += loss
    model.train()
    return epoch_loss / len(val_loader), epoch_acc / len(val_loader)
예제 #6
0
파일: train.py 프로젝트: bysen32/paper00
def train(train_loader, model, criterion_encoder, optimizer_encoder, classifier, criterion_fc, optimizer_fc, epoch, opt):
    model.train()

    losses = AverageMeter()
    fc_losses = AverageMeter()
    encoder_losses = AverageMeter()
    top1 = AverageMeter()

    _print("--" * 50)
    for i, (images, labels) in enumerate(train_loader):
        images = torch.cat(images, dim=0).cuda()
        labels = torch.cat([labels, labels], dim=0).cuda()
        BSZ = labels.shape[0]

        warmup_learning_rate(opt, epoch, i, len(
            train_loader), optimizer_encoder)
        warmup_learning_rate(opt, epoch, i, len(train_loader), optimizer_fc)

        features = model(images, labels)
        encoder_loss = criterion_encoder(features, labels)
        encoder_losses.update(encoder_loss.item(), BSZ)

        logits = classifier(features)
        fc_loss = criterion_fc(logits, labels)
        fc_losses.update(fc_loss.item(), BSZ)

        acc1, acc5 = accuracy(logits, labels, topk=(1, 5))
        top1.update(acc1[0], BSZ)

        loss = encoder_loss + fc_loss
        # loss = fc_loss
        losses.update(loss.item(), BSZ)

        optimizer_encoder.zero_grad()
        optimizer_fc.zero_grad()
        loss.backward()
        optimizer_encoder.step()
        optimizer_fc.step()

        progress_bar(i, len(train_loader), "train")

    _print("epoch:{}".format(epoch))
    _print(
        "loss: {loss.avg:.4f} "
        "encoder_loss: {encoder_loss.avg:.4f} "
        "fc_loss: {fc_loss.avg:.4f} "
        "train acc@1 {top1.avg:.4f} ".format(
            loss=losses, encoder_loss=encoder_losses, fc_loss=fc_losses, top1=top1))
    return losses.avg, top1.avg
예제 #7
0
def test(model, dataset):
    for images, labels in dataset:
        cams = model(images, training=False, using_cam=True)

        cams = tf.image.resize(cams, (image_size, image_size),
                               method='bilinear')
        cams = normalize(cams) * 255

        for image, cam in zip(images, cams.numpy()):
            for class_index in range(classes):
                cv2.imshow('cam - {}'.format(class_index),
                           cam[..., class_index].astype(np.uint8))

            cv2.imshow('show', image.astype(np.uint8))
            cv2.waitKey(0)
예제 #8
0
파일: app.py 프로젝트: gujita/format_ocr
def ocr(absolute_path, img_type, tune_angle=0):
    print("OCR Starting!")
    img = Image.open(absolute_path).convert("RGB")
    t = time.time()
    if tune_angle == 1:
        global_tune = True
        fine_tune = True
    else:
        global_tune = False
        fine_tune = False
    params = select_config(img_type)
    img, res_origin, res_sorted, angle = core.model(img,
                                                    global_tune=global_tune,  # 图片的整体大方向调整,逆时针旋转 镜像. 大约0.5s
                                                    fine_tune=fine_tune,
                                                    # 微调倾斜角(如果能保证图像水平,或者global_tune之后为水平,则不需要微调). 大约2s
                                                    config=params[0],
                                                    if_im=params[1],
                                                    left_adjust=params[2],  # 对检测的文本行进行向左延伸
                                                    right_adjust=params[3],  # 对检测的文本行进行向右延伸
                                                    alpha=params[4],  # 对检测的文本行进行向右、左延伸的倍数
                                                    result_typeset_opotion=0,  # 结果排版方案 TODO: other two type
                                                    )
    print("检测识别1  总耗时:{}s\n".format(time.time() - t))
    json_str = ''
    # for index, _ in enumerate(re_origin):
    #     print(res_origin[index]["text"])
    if len(res_sorted) < 1:
        # TODO: return
        return 'we are sorry, no text detected!'
    else:
        basename = os.path.basename(absolute_path)
        print(basename)
        out_path = './users/out_result/' + basename[:-4] + '.txt'
        print(out_path)
        out = open(out_path, 'w')
        for row in range(res_sorted.shape[0]):
            for col in range(res_sorted.shape[1]):
                json_str += res_sorted[row][col]
                out.write(res_sorted[row][col])
            json_str += '\n'
            out.write('\n')
        out.close()
    print(json_str)
    return json_str
예제 #9
0
def train(train_loader, model, optimizer, criterion):
    model.train()
    epoch_loss = 0
    bar = tqdm(total=len(train_loader))
    b_ix = 1
    for X, Y in train_loader:
        X = X.to(c.device)
        Y = Y.to(c.device)
        optimizer.zero_grad()
        outputs = model(X)
        loss = criterion(outputs, Y)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()

        if b_ix % 10 == 0:
            bar.update(10)
            bar.set_description('current loss:{:.4f}'.format(epoch_loss / b_ix))
        b_ix += 1
    bar.update((b_ix - 1) % 10)
    bar.close()
    return epoch_loss / len(train_loader)
예제 #10
0
파일: train.py 프로젝트: bysen32/paper00
def validate(val_loader, model, criterion_encoder, classifier, criterion_fc, epoch, opt):
    model.eval()

    losses = AverageMeter()
    fc_losses = AverageMeter()
    encoder_losses = AverageMeter()
    top1 = AverageMeter()

    with torch.no_grad():
        for i, (images, labels) in enumerate(val_loader):
            images = images.cuda()
            labels = labels.cuda()
            BSZ = labels.shape[0]

            # caculate class loss
            features = model(images, labels)
            # encoder_loss = criterion_encoder(features, labels)
            # encoder_losses.update(encoder_loss.item(), BSZ)

            logits = classifier(features)
            fc_loss = criterion_fc(logits, labels)
            fc_losses.update(fc_loss.item(), BSZ)

            loss = fc_loss
            losses.update(loss.item(), BSZ)

            acc1, acc5 = accuracy(logits, labels, topk=(1, 5))
            top1.update(acc1[0], BSZ)

            progress_bar(i, len(val_loader), "eval test set")

    _print(
        "loss: {loss.avg:.4f} "
        "fc_loss :{fc_loss.avg:.4f} "
        "acc@1 {top1.avg:.4f}".format(loss=losses, fc_loss=fc_losses, top1=top1))

    return losses.avg, top1.avg
예제 #11
0
SOFTWARE.
"""

__author__ = "Marcus T. Andersson"
__copyright__ = "Copyright 2020, Marcus T. Andersson"
__credits__ = ["Marcus T. Andersson"]
__license__ = "MIT"
__version__ = "2"
__maintainer__ = "Marcus T. Andersson"

import core

statements = []

# Entities have Identity
statements += core.model("EntityE -- LabelEScn1")
statements += core.model("EntityE -- IdentityEScn1")

# Files have Path
statements += core.model("FileE -- PathEScn1")

# Mutables can change content but Constants remain the same
statements += core.model("MutableE -- ContentEEcnn -- ConstantE")
statements += core.model("ConstantE -- ContentTypeEScn1")
statements += core.model("ConstantE -- CreationTimeEScn1")
statements += core.model("ConstantE -- ContentEBcn1")
statements += core.model("ConstantE -- ShaEScn1")

# Embedded things are located in Containers
statements += core.model("ContainerE -- EmbeddedEEcnn -- EmbeddedE")
예제 #12
0
def rolling_window(model_config,
                   datasetName,
                   start_from,
                   restart,
                   restart_window,
                   window_length,
                   timesteps,
                   max_epochs,
                   data_folder,
                   save_folder,
                   col='log_ret',
                   returns=True):
    ### Model:

    model, loss_fn, optimiser = give_me_model(model_config)

    ### CUDA for PyTorch:
    use_cuda = torch.cuda.is_available()

    device = torch.device("cuda:0" if use_cuda else "cpu")
    cudnn.benchmark = True

    model = model
    if torch.cuda.is_available():
        print("We're running on GPU")
        model.cuda()

    ### Model end
    path = data_folder + datasetName + '.csv'
    dataset = miniDataLoader(path, start_from, cols=col, returns=returns)

    train_dt = dataset.get_data(timesteps, False, 1)
    predictions = []
    true_labels = []

    end_of_window = 253 * window_length

    end_of_loop = len(train_dt[0])
    print(end_of_loop)
    i = 0
    while end_of_window + i + 1 < end_of_loop:
        print(i)
        t0 = time.time()
        train_windows = (train_dt[0][i:end_of_window + i],
                         train_dt[1][i:end_of_window + i])
        test_window = (train_dt[0][end_of_window + i + 1],
                       train_dt[1][end_of_window + i + 1])
        train_windows = Dataset(train_windows)
        dataloader_params = {
            'batch_size': 1,
            'shuffle': True,
            'drop_last': True,
            'num_workers': 0
        }
        training_windows_generator = data.DataLoader(train_windows,
                                                     **dataloader_params)

        avg_loss_epoch = np.zeros((max_epochs, 1))
        best_accuracy = torch.FloatTensor([1000])

        if i >= restart_window:
            restart = False

        if restart is False:
            for epoch in range(max_epochs):
                # print("Epoch nr: " + str(epoch))
                loss_vals = np.zeros(
                    (max_epochs, len(training_windows_generator)))
                batch_nr = 0
                for batch, labels in training_windows_generator:
                    model.batch_size = 1
                    batch = batch.view(timesteps, 1, -1)
                    labels = labels.float()

                    # Transfer to GPU
                    batch, labels = batch.to(device), labels.to(device)
                    optimiser.zero_grad()
                    preds = model(batch)
                    loss = loss_fn(preds.view(-1), labels)
                    loss_vals[epoch, batch_nr] = loss.item()
                    loss.backward()
                    optimiser.step()

                    batch_nr += 1
                avg_loss_epoch[epoch] = np.mean(loss_vals[epoch, :])
                is_best = bool(avg_loss_epoch[epoch] < best_accuracy.numpy())

                best_accuracy = torch.FloatTensor(
                    min(avg_loss_epoch[epoch], best_accuracy.numpy()))
                save_checkpoint(
                    {
                        'epoch': 0 + epoch + 1,
                        'state_dict': model.state_dict(),
                        'best_accuracy': best_accuracy
                    }, i, is_best, save_folder + datasetName)
        print(save_folder)
        path_to_checkpoint = save_folder + datasetName + '/' + 'checkpoint' + str(
            i) + '.pth.tar'
        checkpoint = torch.load(path_to_checkpoint)
        model.load_state_dict(checkpoint['state_dict'])
        test_seq = torch.from_numpy(test_window[0]).float().view(
            timesteps, 1, -1).to(device)
        predictions.append(model(test_seq))
        true_labels.append(test_window[1])
        output = (predictions, true_labels, dataset.dates)
        file_output = open(save_folder + datasetName + '/last_pickle.obj',
                           'wb')
        pickle.dump(output, file_output)
        print('The window took {} seconds'.format(time.time() - t0))
        i += 1

    output = (predictions, true_labels)
    return output
예제 #13
0
파일: demo.py 프로젝트: gujita/format_ocr
if __name__ == '__main__':
    print("OCR Starting!")
    img = Image.open(paths[0]).convert("RGB")
    t = time.time()
    # TODO if res_sorted is null
    img, res_origin, res_sorted, angle = core.model(
        img,
        global_tune=False,  # 图片的整体大方向调整,逆时针旋转 镜像. 大约0.5s
        fine_tune=False,
        # 微调倾斜角(如果能保证图像水平,或者global_tune之后为水平,则不需要微调). 大约2s
        config=dict(
            MAX_HORIZONTAL_GAP=50,
            # (1.5~2.5较佳)字符之间的最大间隔,用于文本行的合并 TODO:最好是自动计算
            MIN_V_OVERLAPS=0.6,  # 小 ==》斜 TODO
            MIN_SIZE_SIM=0.6,
            TEXT_PROPOSALS_MIN_SCORE=0.1,
            # 值越小,候选框越多(一些模棱两可的文字)
            TEXT_PROPOSALS_NMS_THRESH=0.3  # 候选框非极大值抑制
        ),
        if_im=False,
        left_adjust=False,  # 对检测的文本行进行向左延伸
        right_adjust=False,  # 对检测的文本行进行向右延伸
        alpha=0.2,  # 对检测的文本行进行向右、左延伸的倍数
        result_typeset_opotion=0,  # 结果排版方案 TODO: other two type
    )
    print("检测识别1  总耗时:{}s\n".format(time.time() - t))
    # for index, _ in enumerate(re_origin):
    #     print(res_origin[index]["text"])
    if len(res_sorted) < 1:
        print("we are sorry, no text are detected!")
    else:
예제 #14
0
 print("OCR Starting!")
 with open('testA_part_result_test.csv', 'w') as csv_file:
     fieldheader = ['filename', 'content']
     writer = csv.DictWriter(csv_file, fieldheader)
     writer.writeheader()
     for path in paths:
         img = Image.open(path).convert("RGB")
         _, result, angle = core.model(
             img,
             global_tune=cfg.global_tune,  # 图片的整体大方向调整,逆时针旋转 镜像. 大约0.5s
             fine_tune=cfg.
             fine_tune,  # 微调倾斜角(如果能保证图像水平,或者global_tune之后为水平,则不需要微调). 大约2s
             config=dict(
                 MAX_HORIZONTAL_GAP=80,  # 字符之间的最大间隔,用于文本行的合并 TODO:最好是自动计算
                 MIN_V_OVERLAPS=0.6,  # 小 ==》斜 TODO
                 MIN_SIZE_SIM=0.6,
                 TEXT_PROPOSALS_MIN_SCORE=0.1,  # 值越小,候选框越多(一些模棱两可的文字)
                 TEXT_PROPOSALS_NMS_THRESH=0.4  # 候选框非极大值抑制
             ),
             if_im=cfg.if_im,
             left_adjust=True,  # 对检测的文本行进行向左延伸
             right_adjust=True,  # 对检测的文本行进行向右延伸
             alpha=0.2  # 对检测的文本行进行向右、左延伸的倍数
         )
         content = ''
         row = {}  # TODO 修改为多行一次写入
         for index, _ in enumerate(result):
             content += result[index]["text"]
         row.update(filename=path[17:-4], content=content)
         writer.writerow(row)
     # writer.writerow({'mp': '7-DForc.郭美a刚g足扇僧:29370777707900273970ee0.o)110193721:2997311国警', 'mp1': '合号单币鄂昨eo电舒y营柱国)警翠女义职翠餐武署益罩发兴'})
def train(args):
    """Sets up the model to train"""
    # Create a writer object to log events during training
    writer = SummaryWriter(pjoin('runs', 'exp_1'))

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Load splits
    x_train, y_train, x_val, y_val, seismic = train_val_split(args)

    # Convert to torch tensors in the form (N, C, L)
    x_train = torch.from_numpy(np.expand_dims(x_train, 1)).float().to(device)
    y_train = torch.from_numpy(np.expand_dims(y_train, 1)).float().to(device)
    x_val = torch.from_numpy(np.expand_dims(x_val, 1)).float().to(device)
    y_val = torch.from_numpy(np.expand_dims(y_val, 1)).float().to(device)
    seismic = torch.from_numpy(np.expand_dims(seismic, 1)).float().to(device)

    # Set up the dataloader for training dataset
    dataset = SeismicLoader(x_train, y_train)
    train_loader = DataLoader(dataset=dataset,
                              batch_size=args.batch_size,
                              shuffle=False)

    # import tcn
    model = TCN(1, 1, args.tcn_layer_channels, args.kernel_size,
                args.dropout).to(device)

    #model = ANN().to(device)

    #model = lstm().to(device)

    # Set up loss
    criterion = torch.nn.MSELoss()

    # Define Optimizer
    optimizer = torch.optim.Adam(model.parameters(),
                                 weight_decay=args.weight_decay,
                                 lr=args.lr)

    # Set up list to store the losses
    train_loss = [np.inf]
    val_loss = [np.inf]
    iter = 0
    # Start training
    for epoch in range(args.n_epoch):
        for x, y in train_loader:
            model.train()
            optimizer.zero_grad()
            y_pred = model(x)
            loss = criterion(y_pred, y)
            loss.backward()
            optimizer.step()
            train_loss.append(loss.item())
            writer.add_scalar(tag='Training Loss',
                              scalar_value=loss.item(),
                              global_step=iter)
            if epoch % 200 == 0:
                with torch.no_grad():
                    model.eval()
                    y_pred = model(x_val)
                    loss = criterion(y_pred, y_val)
                    val_loss.append(loss.item())
                    writer.add_scalar(tag='Validation Loss',
                                      scalar_value=loss.item(),
                                      global_step=iter)
            print(
                'epoch:{} - Training loss: {:0.4f} | Validation loss: {:0.4f}'.
                format(epoch, train_loss[-1], val_loss[-1]))

            if epoch % 100 == 0:
                with torch.no_grad():
                    model.eval()
                    AI_inv = model(seismic)
                fig, ax = plt.subplots()
                ax.imshow(AI_inv[:, 0].detach().cpu().numpy().squeeze().T,
                          cmap="rainbow")
                ax.set_aspect(4)
                writer.add_figure('Inverted Acoustic Impedance', fig, iter)
        iter += 1

    writer.close()

    # Set up directory to save results
    results_directory = 'results'
    seismic_offsets = np.expand_dims(marmousi_seismic().squeeze()[:, 100:600],
                                     1)
    seismic_offsets = torch.from_numpy(
        (seismic_offsets - seismic_offsets.mean()) /
        seismic_offsets.std()).float()
    with torch.no_grad():
        model.cpu()
        model.eval()
        AI_inv = model(seismic_offsets)

    if not os.path.exists(
            results_directory
    ):  # Make results directory if it doesn't already exist
        os.mkdir(results_directory)
        print('Saving results...')
    else:
        print('Saving results...')

    np.save(pjoin(results_directory, 'AI.npy'),
            marmousi_model().T[452:2399, 400:2400])
    np.save(pjoin(results_directory, 'AI_inv.npy'),
            AI_inv.detach().numpy().squeeze()[452:2399])
    print('Results successfully saved.')
예제 #16
0
파일: training.py 프로젝트: stas97/bsc_lstm
def train_model(model, loss, optimiser, scheduler, max_epochs, train_gen, test_gen,
                timesteps, batch_size, folder, resuming):
    ######## CUDA for PyTorch
    use_cuda = torch.cuda.is_available()

    device = torch.device("cuda:0" if use_cuda else "cpu")
    cudnn.benchmark = True

    model = model
    if torch.cuda.is_available():
        print("We're running on GPU")
        model.cuda()
    #######
    scheduler = scheduler
    optimiser = optimiser
    loss_fn = loss

    # Track losses:

    loss_vals_train = np.zeros((max_epochs, len(train_gen)))
    loss_vals_test = np.zeros((max_epochs, len(test_gen)))

    avg_loss_epoch_train = np.zeros((max_epochs, 1))
    avg_loss_epoch_test = np.zeros((max_epochs, 1))

    print("Start Training")
    total_time = 0
    best_accuracy = torch.FloatTensor([1000])

    start_epoch = 0
    if resuming:
        path_to_checkpoint = folder + '/checkpoint.pth.tar'
        cuda = torch.cuda.is_available()
        if cuda:
            checkpoint = torch.load(path_to_checkpoint)
        else:
            # Load GPU model on CPU
            checkpoint = torch.load(path_to_checkpoint,
                                    map_location=lambda storage,
                                                        loc: storage)
        start_epoch = checkpoint['epoch']
        best_accuracy = checkpoint['best_accuracy']
        model.load_state_dict(checkpoint['state_dict'])
        print("=> loaded checkpoint '{}' (trained for {} epochs)".format(path_to_checkpoint, checkpoint['epoch']))

    for epoch in range(start_epoch, max_epochs, 1):
        print("Epoch nr: " + str(epoch))
        t0 = time.time()
        batch_nr = 0
        for batch, labels in train_gen:
            model.batch_size = batch_size
            batch = batch.view(timesteps, batch_size, -1)
            # print(batch.shape)
            labels = labels.float()
            # Transfer to GPU
            batch, labels = batch.to(device), labels.to(device)
            optimiser.zero_grad()

            preds = model(batch)
            loss = loss_fn(preds, labels)
            loss_vals_train[epoch, batch_nr] = loss.item()

            loss.backward()

            optimiser.step()

            batch_nr += 1
        if scheduler is not None:
            scheduler.step()

        batch_nr = 0
        for batch, labels in test_gen:
            model.batch_size = 1
            batch = batch.view(timesteps, 1, -1)
            labels = labels.float()
            # Transfer to GPU
            batch, labels = batch.to(device), labels.to(device)
            optimiser.zero_grad()

            y_pred_test = model(batch)
            loss = loss_fn(y_pred_test, labels)
            loss_vals_test[epoch, batch_nr] = loss.item()
            batch_nr += 1

        error = np.mean(loss_vals_test[epoch, :])
        print('The epoch took {} seconds'.format(time.time() - t0))
        total_time += time.time() - t0

        avg_loss_epoch_train[epoch] = np.mean(loss_vals_train[epoch, :])
        avg_loss_epoch_test[epoch] = np.mean(loss_vals_test[epoch, :])

        is_best = bool(avg_loss_epoch_test[epoch] < best_accuracy.numpy())

        best_accuracy = torch.FloatTensor(min(avg_loss_epoch_test[epoch], best_accuracy.numpy()))
        print(best_accuracy)
        save_checkpoint({
            'epoch': start_epoch + epoch + 1,
            'state_dict': model.state_dict(),
            'best_accuracy': best_accuracy
        }, is_best, folder)

    print('Total training time is: ' + str(total_time) + ' seconds')
    save_path = folder + '/' + 'last_model.pth.tar'
    torch.save({
        'epoch': start_epoch + epoch + 1,
        'state_dict': model.state_dict(),
        'best_accuracy': best_accuracy
    }, save_path)

    with open(folder + '/model_summary.txt', 'w+') as f:
        f.write(str(model))  # Python 3.x
    print("Last Model Saved")
    return model, avg_loss_epoch_train, avg_loss_epoch_test, error