def train(loaders, save_path): """returns trained model""" # Initialize custom defined cnn model = Net() use_cuda = torch.cuda.is_available() if use_cuda: model.cuda() # cross entropy loss for classification task criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr=config.lr) # initialize tracker for minimum validation loss valid_loss_min = np.Inf n_epochs = config.n_epochs for epoch in range(1, n_epochs + 1): # initialize variables to monitor training and validation loss train_loss = 0.0 valid_loss = 0.0 model.train() for batch_idx, (data, target) in enumerate(loaders['train']): # move to GPU if use_cuda: data, target = data.cuda(), target.cuda() optimizer.zero_grad() output = model(data) loss = criterion(output, target) loss.backward() optimizer.step() # average training loss train_loss += (1 / (batch_idx + 1)) * (loss.data - train_loss) # vaidation model.eval() for batch_idx, (data, target) in enumerate(loaders['valid']): # move to GPU if use_cuda: data, target = data.cuda(), target.cuda() ## update the average validation loss output = model(data) loss = criterion(output, target) valid_loss += (1 / (batch_idx + 1)) * (loss.data - valid_loss) # print training/validation statistics print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f}'. format(epoch, train_loss, valid_loss)) # save the model if validation loss has decreased if valid_loss <= valid_loss_min: torch.save(model.state_dict(), save_path) # Updating the validation loss minimum valid_loss_min = valid_loss # return trained model return model
args.cnn_helmet_resize_width, args.cnn_helmet_resize_height)), transforms.ToTensor()]) train_data = datasets.ImageFolder(args.dataset_location, transform=train_transforms) train_loader = data_utils.DataLoader(train_data, batch_size=args.batch_size, shuffle=True, drop_last = True) model = Net() if torch.cuda.is_available(): model.cuda() criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters(), lr = args.learning_rate) for epoch in range(args.epoch): start_time = time.time() running_loss = 0.0 for i, data in enumerate(train_loader, 0): inputs, labels = data if torch.cuda.is_available(): inputs = inputs.cuda() labels = labels.cuda() optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels)