Ejemplo n.º 1
0
def main():

    print(1)
    args = parser.parse_args()

    torch.manual_seed(999)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    D_xs = Net(args.window_size)
    D_xs.apply(weights_init)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    #torch.cuda.set_device(0)
    x = torch.FloatTensor(args.batch_size, 3, args.window_size, 256, 256)

    args.cuda = True
    if args.cuda:
        print("Let's use", torch.cuda.device_count(), "GPUs!")
        D_xs = torch.nn.DataParallel(D_xs).cuda()
        #D_xs = D_xs.cuda()

        x = x.cuda()
    x = Variable(x)
    #D_xs.to(device)

    print(2)

    lr = args.learning_rate
    lr_s = args.learning_rate_s
    ourBetas = [args.momentum, args.momentum2]
    batch_size = args.batch_size
    snapshot = args.nsnapshot

    D_xs_solver = optim.Adam(
        D_xs.parameters(),
        lr=lr,
    )

    cudnn.benchmark = True
    l1Loss = nn.L1Loss().cuda()
    e_shift = 0
    min_val_loss = 99999
    no_improve_epoch = 0
    now = datetime.now()
    print('begin')

    log_path = 'log/lr_{}_time:{}'.format(args.learning_rate,
                                          now.strftime("%Y%m%d-%H%M%S"))
    # try:
    # 	os.mkdir(log_path)
    # except:
    # 	pass
    writer = SummaryWriter(log_path)

    for epoch in range(args.epochs):

        path = 'dataset/'

        train_loader = torch.utils.data.DataLoader(EyeBrowDataset(
            'dataset/train_all.mat',
            '',
            args.window_size,
            transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
            ])),
                                                   batch_size=args.batch_size,
                                                   shuffle=True,
                                                   num_workers=args.workers,
                                                   pin_memory=True)

        print('len of train_loader', len(train_loader))

        mes_sum = 0
        for i, (data, value) in enumerate(train_loader):

            if len(data) < args.batch_size:
                continue

            D_xs.zero_grad()

            imgseq = data

            x.data.resize_(imgseq.size()).copy_(imgseq)
            vv = value.type(torch.FloatTensor)
            vv = vv.cuda()
            #vv = vv.to(device)
            vv = Variable(vv, requires_grad=False)

            #x= x.to(device)

            score = D_xs(x)

            v_loss = l1Loss(score, vv)
            v_loss.backward()
            D_xs_solver.step()
            mes_sum += v_loss.item()

            if True:
                # s.append(mes_sum)
                # scio.savemat('log/log_after'+str(i)+'_iteration.mat',{'loss':s})

                writer.add_scalar('train/loss', mes_sum, i)

                mes_sum = 0

            print('epoch:[%2d] [%4d/%4d] loss: %.4f' %
                  (epoch + e_shift, i, len(data), v_loss.item()))

        print('epoch:', epoch)
        path = 'dataset/'

        vali_loader = torch.utils.data.DataLoader(EyeBrowDataset(
            'dataset/validation_all.mat',
            '',
            args.window_size,
            transform=transforms.Compose([
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
            ])),
                                                  batch_size=args.batch_size,
                                                  shuffle=True,
                                                  num_workers=args.workers,
                                                  pin_memory=True)
        mse_sum = 0

        print('len of train_loader', len(vali_loader))

        for i, (data, value) in enumerate(vali_loader):

            imgseq = data
            with torch.no_grad():
                x.data.resize_(imgseq.size()).copy_(imgseq)
            vv = value.type(torch.FloatTensor)
            vv = vv.cuda()
            vv = Variable(vv, requires_grad=False)

            score = D_xs(x)

            v_loss = l1Loss(score, vv)

            mse_sum = mse_sum + v_loss.item()

            print('epoch:[%2d] [%4d/%4d] loss: %.4f' %
                  (epoch + e_shift, i, len(data), v_loss.item()))

        val_loss = mse_sum / float(i + 1)
        print("*** Epoch: [%2d], "
              "val_mse: %.6f ***" % (epoch + e_shift, val_loss))

        # if performance improve save the new model
        # if performance does not increase for patiences epochs, stop training
        if val_loss < min_val_loss:
            min_val_loss = val_loss
            no_improve_epoch = 0
            val_loss = round(val_loss, 2)
            torch.save(
                D_xs.state_dict(),
                '{}/netD_xs_epoch_{}_val_loss_{}.pth'.format(
                    args.outf, epoch + e_shift, val_loss))
            print("performance improve, saved the new model......")
        else:
            no_improve_epoch += 1

        if no_improve_epoch > args.patiences:
            print("stop training....")
            break
Ejemplo n.º 2
0
def train(n_epochs=n_epochs,
          train_loader=train_loader,
          valid_loader=valid_loader,
          save_location_path=save_location_path):

    train_on_gpu = torch.cuda.is_available()

    def init_weights(m):
        if isinstance(m, nn.Linear):
            I.xavier_uniform_(m.weight)

    model = Net()
    model.apply(init_weights)

    if train_on_gpu:
        model.cuda()

    criterion = nn.MSELoss()
    optimizer = optim.Adam(params=model.parameters(), lr=0.001)

    valid_loss_min = np.Inf

    model.train()

    for epoch in range(1, n_epochs + 1):
        # Keep track of training and validation loss
        train_loss = 0.0
        valid_loss = 0.0

        for data in train_loader:

            # Grab the image and its corresponding label
            images = data['image']
            key_pts = data['keypoints']

            if train_on_gpu:
                images, key_pts = images.cuda(), key_pts.cuda()

            # Flatten keypoints & convert data to FloatTensor for regression loss
            key_pts = key_pts.view(key_pts.size(0), -1)
            if train_on_gpu:
                key_pts = key_pts.type(torch.cuda.FloatTensor)
                images = images.type(torch.cuda.FloatTensor)
            else:
                key_pts = key_pts.type(torch.FloatTensor)
                images = images.type(torch.FloatTensor)

            optimizer.zero_grad()  # Clear the gradient
            output = model(images)  # Forward
            loss = criterion(output, key_pts)  # Compute the loss
            loss.backward()  # Compute the gradient
            optimizer.step()  # Perform updates using calculated gradients

            train_loss += loss.item() * images.size(0)

        # Validation
        model.eval()
        for data in valid_loader:

            images = data['image']
            key_pts = data['keypoints']

            if train_on_gpu:
                images, key_pts = images.cuda(), key_pts.cuda()

            key_pts = key_pts.view(key_pts.size(0), -1)
            if train_on_gpu:
                key_pts = key_pts.type(torch.cuda.FloatTensor)
                images = images.type(torch.cuda.FloatTensor)
            else:
                key_pts = key_pts.type(torch.FloatTensor)
                images = images.type(torch.FloatTensor)

            output = model(images)
            loss = criterion(output, key_pts)

            valid_loss += loss.item() * images.size(0)

        # calculate average losses
        train_loss = train_loss / len(train_loader)
        valid_loss = valid_loss / len(valid_loader)

        print(
            f"epoch: {epoch} \t trainLoss: {train_loss} \t valLoss: {valid_loss}"
        )

        eviz.send_data(current_epoch=epoch,
                       current_train_loss=train_loss,
                       current_val_loss=valid_loss)
Ejemplo n.º 3
0
class Training():
    def __init__(self):
        self.net = Net()
        self.net.apply(self.init_weights)
        #Basic logging
        logging.basicConfig(filename="cnn2.log", level=logging.DEBUG)
        logging.info(self.net)
        logging.info("Number of parameters: {}".format(
            self.count_parameters(self.net)))

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

        self.optimizer = torch.optim.SGD(self.net.parameters(),
                                         lr=args.lr,
                                         momentum=0.9)
        self.criterion = nn.CrossEntropyLoss().to(self.device)

        self.best_acc = 0
        self.net.to(self.device)

    def loader(self):

        # Define transformers to apply on the input data

        transform_train = transforms.Compose([
            transforms.RandomCrop(28, padding=4),
            transforms.ToTensor(),
            # transforms.Normalize((mean,), (std,)),
        ])

        transform_valid = transforms.Compose([
            transforms.ToTensor(),
            # transforms.Normalize((mean,), (std,)),
        ])

        train = datasets.EMNIST(args.root,
                                split="balanced",
                                train=True,
                                download=True,
                                transform=transform_train)
        test = datasets.EMNIST(args.root,
                               split="balanced",
                               train=False,
                               download=True,
                               transform=transform_valid)

        self.train_loader = torch.utils.data.DataLoader(
            train,
            batch_size=args.batch_size,
            shuffle=True,
            num_workers=2,
            drop_last=True)

        self.test_loader = torch.utils.data.DataLoader(
            test,
            batch_size=args.test_batch_size,
            shuffle=False,
            num_workers=2,
            drop_last=True)

    def init_weights(self, m):
        if type(m) == nn.Linear:
            torch.nn.init.xavier_uniform_(m.weight)
            m.bias.data.fill_(0.01)

    def count_parameters(self, model):
        return sum(p.numel() for p in model.parameters() if p.requires_grad)

    def inf_generator(self, iterable):
        """Allows training with DataLoaders in a single infinite loop:
			for i, (x, y) in enumerate(inf_generator(train_loader)):
		"""
        iterator = iterable.__iter__()
        while True:
            try:
                yield iterator.__next__()
            except StopIteration:
                iterator = iterable.__iter__()

    def train(self):
        self.loader()

        data_gen = self.inf_generator(self.train_loader)
        batches_per_epoch = len(self.train_loader)
        for itr in range(args.EPOCHS * batches_per_epoch):

            self.optimizer.zero_grad()
            x, y = data_gen.__next__()
            x = x.view(-1, 28, 28, 1)
            x = torch.transpose(x, 1, 2)

            x = x.to(self.device)
            y = y.to(self.device)
            logits = self.net(x)
            loss = self.criterion(logits, y)

            loss.backward()
            self.optimizer.step()

            if itr % batches_per_epoch == 0:
                with torch.no_grad():
                    train_acc = self.accuracy(self.net, self.train_loader)
                    val_acc = self.accuracy(self.net, self.test_loader)
                    if val_acc > self.best_acc:
                        torch.save({"state_dict": self.net.state_dict()},
                                   "alpha_weights.pth")
                        self.best_acc = val_acc
                    logging.info("Epoch {:04d}"
                                 "Train Acc {:.4f} | Test Acc {:.4f}".format(
                                     itr // batches_per_epoch, train_acc,
                                     val_acc))

                    print("Epoch {:04d}"
                          "Train Acc {:.4f} | Test Acc {:.4f}".format(
                              itr // batches_per_epoch, train_acc, val_acc))

    def one_hot(self, x, K):
        return np.array(x[:, None] == np.arange(K)[None, :], dtype=int)

    def accuracy(self, model, dataset_loader):
        total_correct = 0
        for x, y in dataset_loader:
            x = x.view(-1, 28, 28, 1)
            x = torch.transpose(x, 1, 2)

            x = x.to(self.device)
            y = self.one_hot(np.array(y.numpy()), 47)

            target_class = np.argmax(y, axis=1)
            predicted_class = np.argmax(model(x).cpu().detach().numpy(),
                                        axis=1)
            total_correct += np.sum(predicted_class == target_class)
        return total_correct / len(dataset_loader.dataset)