本文整理汇总了Python中tensorboardX.SummaryWriter.add_text方法的典型用法代码示例。如果您正苦于以下问题：Python SummaryWriter.add_text方法的具体用法？Python SummaryWriter.add_text怎么用？Python SummaryWriter.add_text使用的例子？那么恭喜您, 这里精选的方法代码示例或许可以为您提供帮助。您也可以进一步了解该方法所在类tensorboardX.SummaryWriter的用法示例。
在下文中一共展示了SummaryWriter.add_text方法的1个代码示例，这些例子默认根据受欢迎程度排序。您可以为喜欢或者感觉有用的代码点赞，您的评价将有助于系统推荐出更棒的Python代码示例。
示例1: __init__
# 需要导入模块: from tensorboardX import SummaryWriter [as 别名]
# 或者: from tensorboardX.SummaryWriter import add_text [as 别名]
class Train:
__device = []
__writer = []
__model = []
__transformations = []
__dataset_train = []
__train_loader = []
__loss_func = []
__optimizer = []
__exp_lr_scheduler = []
def __init__(self, gpu='0'):
# Device configuration
self.__device = torch.device('cuda:'+gpu if torch.cuda.is_available() else 'cpu')
self.__writer = SummaryWriter('logs')
self.__model = CNNDriver()
# Set model to train mode
self.__model.train()
print(self.__model)
self.__writer.add_graph(self.__model, torch.rand(10, 3, 66, 200))
# Put model on GPU
self.__model = self.__model.to(self.__device)
def train(self, num_epochs=100, batch_size=400, lr=0.0001, l2_norm=0.001, save_dir='./save', input='./DataLMDB'):
# Create log/save directory if it does not exist
if not os.path.exists('./logs'):
os.makedirs('./logs')
if not os.path.exists(save_dir):
os.makedirs(save_dir)
self.__transformations = transforms.Compose([AugmentDrivingTransform(),
RandomBrightness(), ConvertToGray(),
ConvertToSepia(), AddNoise(), DrivingDataToTensor(),])
self.__dataset_train = DriveData_LMDB(input, self.__transformations)
self.__train_loader = DataLoader(self.__dataset_train, batch_size=batch_size, shuffle=True, num_workers=4)
# Loss and Optimizer
self.__loss_func = nn.MSELoss()
# self.__loss_func = nn.SmoothL1Loss()
self.__optimizer = torch.optim.Adam(self.__model.parameters(), lr=lr, weight_decay=l2_norm)
# Decay LR by a factor of 0.1 every 10 epochs
self.__exp_lr_scheduler = lr_scheduler.StepLR(self.__optimizer, step_size=15, gamma=0.1)
print('Train size:', len(self.__dataset_train), 'Batch size:', batch_size)
print('Batches per epoch:', len(self.__dataset_train) // batch_size)
# Train the Model
iteration_count = 0
for epoch in range(num_epochs):
for batch_idx, samples in enumerate(self.__train_loader):
# Send inputs/labels to GPU
images = samples['image'].to(self.__device)
labels = samples['label'].to(self.__device)
self.__optimizer.zero_grad()
# Forward + Backward + Optimize
outputs = self.__model(images)
loss = self.__loss_func(outputs, labels.unsqueeze(dim=1))
loss.backward()
self.__optimizer.step()
self.__exp_lr_scheduler.step(epoch)
# Send loss to tensorboard
self.__writer.add_scalar('loss/', loss.item(), iteration_count)
self.__writer.add_histogram('steering_out', outputs.clone().detach().cpu().numpy(), iteration_count, bins='doane')
self.__writer.add_histogram('steering_in',
labels.unsqueeze(dim=1).clone().detach().cpu().numpy(), iteration_count, bins='doane')
# Get current learning rate (To display on Tensorboard)
for param_group in self.__optimizer.param_groups:
curr_learning_rate = param_group['lr']
self.__writer.add_scalar('learning_rate/', curr_learning_rate, iteration_count)
# Display on each epoch
if batch_idx == 0:
# Send image to tensorboard
self.__writer.add_image('Image', images, epoch)
self.__writer.add_text('Steering', 'Steering:' + str(outputs[batch_idx].item()), epoch)
# Print Epoch and loss
print('Epoch [%d/%d] Loss: %.4f' % (epoch + 1, num_epochs, loss.item()))
# Save the Trained Model parameters
torch.save(self.__model.state_dict(), save_dir+'/cnn_' + str(epoch) + '.pkl')
iteration_count += 1