def __init__(self, cuda, model, optimizer, loss_fcn, scheduler, train_loader, val_loader, out, max_iter, logFile, size_average=False, interval_validate=None): """ :param cuda: :param model: :param optimizer: scheduler:学习率调整策略 :param loss_fcn: :param train_loader: :param val_loader: :param out: 字符串,模型输出的路径,用于保存模型 :param max_iter: :param size_average: :param interval_validate: """ self.cuda = cuda self.model = model self.optim = optimizer self.scheduler = scheduler self.loss_fcn = loss_fcn self.train_loader = train_loader self.val_loader = val_loader self.out = out self.epoch = 0 self.iteration = 0 self.max_iter = max_iter self.best_mean_iu = 0 self.viusal = utils.Visualizer() self.valid_loss = 0 self.valid_acc = 0 self.valMeanIu = 0 self.train_loss = 0 self.train_acc = 0 self.trainMeanIu = 0 self.best_mean_iu = 0 self.logFile = logFile if interval_validate is None: self.interval_validate = len(self.train_loader) else: self.interval_validate = interval_validate
def train(**kwargs): for k_, v_ in kwargs.items(): setattr(opt, k_, v_) vis = utils.Visualizer(opt.env) # 数据加载 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)) ]) loader = get_loader(batch_size=1, data_path=opt.data_path, img_shape=opt.img_shape, transform=transform) # 转换网络 transformer = TransformerNet().cuda() # transformer.load_state_dict(t.load(opt.model_path, )) #if opt.model_path: # transformer.load_state_dict(t.load(opt.model_path,map_location=lambda _s, _: _s)) # 损失网络 Vgg16 vgg = Vgg19().eval() depthnet = HourGlass().eval() depthnet.load_state_dict(t.load(opt.depth_path)) # print(vgg) # BASNET net = BASNet(3, 1).cuda() net.load_state_dict(torch.load('./basnet.pth')) net.eval() # 优化器 optimizer = t.optim.Adam(transformer.parameters(), lr=opt.lr) # 获取风格图片的数据 img = Image.open(opt.style_path) img = img.resize(opt.img_shape) img = transform(img).float() style = Variable(img, requires_grad=True).unsqueeze(0) vis.img('style', (style[0] * 0.225 + 0.45).clamp(min=0, max=1)) if opt.use_gpu: transformer.cuda() style = style.cuda() vgg.cuda() depthnet.cuda() # 风格图片的gram矩阵 style_v = Variable(style, volatile=True) features_style = vgg(style_v) gram_style = [Variable(utils.gram_matrix(y.data)) for y in features_style] # 损失统计 style_meter = tnt.meter.AverageValueMeter() content_meter = tnt.meter.AverageValueMeter() temporal_meter = tnt.meter.AverageValueMeter() long_temporal_meter = tnt.meter.AverageValueMeter() depth_meter = tnt.meter.AverageValueMeter() # tv_meter = tnt.meter.AverageValueMeter() kk = 0 for count in range(opt.epoch): print('Training Start!!') content_meter.reset() style_meter.reset() temporal_meter.reset() long_temporal_meter.reset() depth_meter.reset() # tv_meter.reset() for step, frames in enumerate(loader): for i in tqdm.tqdm(range(1, len(frames))): kk += 1 if (kk + 1) % 3000 == 0: print('LR had changed') for param in optimizer.param_groups: param['lr'] = max(param['lr'] / 1.2, 1e-4) optimizer.zero_grad() x_t = frames[i].cuda() x_t1 = frames[i - 1].cuda() h_xt = transformer(x_t) h_xt1 = transformer(x_t1) depth_x_t = depthnet(x_t) depth_x_t1 = depthnet(x_t1) depth_h_xt = depthnet(h_xt) depth_h_xt1 = depthnet(h_xt1) img1 = h_xt1.data.cpu().squeeze(0).numpy().transpose(1, 2, 0) img2 = h_xt.data.cpu().squeeze(0).numpy().transpose(1, 2, 0) flow, mask = opticalflow(img1, img2) d1, d2, d3, d4, d5, d6, d7, d8 = net(x_t) a1pha1 = PROCESS(d1, x_t) del d1, d2, d3, d4, d5, d6, d7, d8 d1, d2, d3, d4, d5, d6, d7, d8 = net(x_t1) a1pha2 = PROCESS(d1, x_t1) del d1, d2, d3, d4, d5, d6, d7, d8 h_xt_features = vgg(h_xt) h_xt1_features = vgg(h_xt1) x_xt_features = vgg(a1pha1) x_xt1_features = vgg(a1pha2) # ContentLoss, conv3_2 content_t = F.mse_loss(x_xt_features[2], h_xt_features[2]) content_t1 = F.mse_loss(x_xt1_features[2], h_xt1_features[2]) content_loss = opt.content_weight * (content_t1 + content_t) # StyleLoss style_t = 0 style_t1 = 0 for ft_y, gm_s in zip(h_xt_features, gram_style): gram_y = gram_matrix(ft_y) style_t += F.mse_loss(gram_y, gm_s.expand_as(gram_y)) for ft_y, gm_s in zip(h_xt1_features, gram_style): gram_y = gram_matrix(ft_y) style_t1 += F.mse_loss(gram_y, gm_s.expand_as(gram_y)) style_loss = opt.style_weight * (style_t1 + style_t) # # depth loss depth_loss1 = F.mse_loss(depth_h_xt, depth_x_t) depth_loss2 = F.mse_loss(depth_h_xt1, depth_x_t1) depth_loss = opt.depth_weight * (depth_loss1 + depth_loss2) # # TVLoss # print(type(s_hxt[layer]),s_hxt[layer].size()) # tv_loss = TVLoss(h_xt) #Long-temprol loss if (i - 1) % opt.sample_frames == 0: frames0 = h_xt1.cpu() long_img1 = frames0.data.cpu().squeeze( 0).numpy().transpose(1, 2, 0) # long_img2 = h_xt.data.cpu().squeeze(0).numpy().transpose(1,2,0) long_flow, long_mask = opticalflow(long_img1, img2) # Optical flow flow = torch.from_numpy(flow).permute(2, 0, 1).unsqueeze(0).to( torch.float32) long_flow = torch.from_numpy(long_flow).permute( 2, 0, 1).unsqueeze(0).to(torch.float32) # print(flow.size()) # print(h_xt1.size()) warped = warp(h_xt1.cpu().permute(0, 2, 3, 1), flow, opt.img_shape[1], opt.img_shape[0]).cuda() long_warped = warp(frames0.cpu().permute(0, 2, 3, 1), long_flow, opt.img_shape[1], opt.img_shape[0]).cuda() long_temporal_loss = F.mse_loss( h_xt, long_mask * long_warped.permute(0, 3, 1, 2)) # print(warped.size()) # tv.utils.save_image((warped.permute(0,3,1,2).data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1), # './warped.jpg') mask = mask.transpose(2, 0, 1) mask = torch.from_numpy(mask).cuda().to(torch.float32) # print(mask.shape) temporal_loss = F.mse_loss(h_xt, mask * warped.permute(0, 3, 1, 2)) temporal_loss = opt.temporal_weight * temporal_loss long_temporal_loss = opt.long_temporal_weight * long_temporal_loss # Spatial Loss spatial_loss = content_loss + style_loss Loss = spatial_loss + depth_loss + temporal_loss + long_temporal_loss Loss.backward(retain_graph=True) optimizer.step() content_meter.add(float(content_loss.data)) style_meter.add(float(style_loss.data)) temporal_meter.add(float(temporal_loss.data)) long_temporal_meter.add(float(long_temporal_loss.data)) depth_meter.add(float(depth_loss.data)) # tv_meter.add(float(tv_loss.data)) vis.plot('temporal_loss', temporal_meter.value()[0]) vis.plot('long_temporal_loss', long_temporal_meter.value()[0]) vis.plot('content_loss', content_meter.value()[0]) vis.plot('style_loss', style_meter.value()[0]) vis.plot('depth_loss', depth_meter.value()[0]) # vis.plot('tv_loss', tv_meter.value()[0]) if i % 10 == 0: vis.img('input(t)', (x_t.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) vis.img('output(t)', (h_xt.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) vis.img('output(t-1)', (h_xt1.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) print( 'epoch{},content loss:{},style loss:{},temporal loss:{},long temporal loss:{},depth loss:{},total loss{}' .format(count, content_loss, style_loss, temporal_loss, long_temporal_loss, depth_loss, Loss)) # print('epoch{},content loss:{},style loss:{},depth loss:{},total loss{}' # .format(count,content_loss, style_loss,depth_loss,Loss)) vis.save([opt.env]) torch.save(transformer.state_dict(), opt.model_path)
def train(**kwargs): ''' para : opts:the para from your return: the train model ''' opts.parse_kwargs(**kwargs) print "train begin!" viz = utils.Visualizer(opts.env) #model our_model = getattr(models, opts.model)(opts) our_model.load_state_dict( torch.load( "./check_point/<class 'models.RDN.rdn'>_0823_17:59:06.path")) #step2 data_set data_size, data_loader = data.dataset.data_loader(opts) #step 3 criterion optimer #l1 criterion = getattr(models, opts.loss_function)() print data_size optimer = torch.optim.Adam(our_model.parameters(), lr=opts.lr) sche = torch.optim.lr_scheduler.MultiStepLR(optimer, milestones=opts.stone, gamma=0.5) #step 4 device best_loss = 1e11 device = torch.device("cuda:0" if opts.use_gpu else "cpu") our_model.to(device) since = time.time() plot_line_win = None plot_img_win = None plot_label_win = None plot_test_img_win = None plot_test_label_win = None #step 5 trainning for epoch in range(10, opts.max_epoch): print("*****************" * 10) print("epoch {}/{}".format(epoch, opts.max_epoch)) our_model.train() epoch_loss = 0.0 for _, datas in enumerate(data_loader['train'], 1): inputs = datas['input'].to(device) labels = datas['label'].to(device) optimer.zero_grad() with torch.set_grad_enabled(True): outputs = our_model(inputs) loss = criterion(outputs, labels) loss.backward() optimer.step() show_outputs = copy.deepcopy(outputs.detach()) show_outputs[show_outputs > 1.0] = 1.0 show_outputs[show_outputs < 0.0] = 0.0 plot_img_win = viz.images(show_outputs, win=plot_img_win, title='img_test') plot_label_win = viz.images(labels, win=plot_label_win, title='img_label') epoch_loss += loss.item() * inputs.shape[0] if _ % 40 == 0: #draw loss_line print "iteration :{}".format(_) times = time.time() - since times = 1.0 * times / 60 x = times train_loss = (epoch_loss * 1.0 / (opts.batch_size * _)) plot_test_img_win, plot_test_label_win, test_loss = val( our_model, data_loader['val'], data_size['val'], device, criterion, plot_test_img_win, plot_test_label_win, viz) x = np.column_stack((np.asarray(x), np.asarray(x))) y = np.column_stack( (np.asarray(train_loss), np.asarray(test_loss))) plot_line_win = viz.plot(x, y, plot_line_win, 'Loss', 'train_loss', 'val_loss') if (test_loss < best_loss): best_loss = test_loss best_model = our_model.state_dict() our_model.save() epoch_loss = epoch_loss / data_size['train'] plot_test_img_win, plot_test_label_win, test_loss = val( our_model, data_loader['val'], data_size['val'], device, criterion, plot_test_img_win, plot_test_label_win, viz) print("{} :train_loss{:.8f},val_loss{:.8f}".format( 'lossing', epoch_loss, test_loss)) sche.step() print("best_loss for val{:.8f}".format(best_loss)) torch.save(opts.load_model_path, best_model)
def train(**kwargs): opt = Config() for k_, v_ in kwargs.items(): setattr(opt, k_, v_) # 可视化操作 vis = utils.Visualizer(opt.env) # 数据加载 transfroms = tv.transforms.Compose([ # 将输入的`PIL.Image`重新改变大小成给定的`size` `size`是最小边的边长 tv.transforms.Scale(opt.image_size), tv.transforms.CenterCrop(opt.image_size), # 转为0-1之间 tv.transforms.ToTensor(), # 转为0-255之间 tv.transforms.Lambda(lambda x: x * 255) ]) # 封装数据集,并进行数据转化 dataset = tv.datasets.ImageFolder(opt.data_root, transfroms) # 数据加载器 dataloader = data.DataLoader(dataset, opt.batch_size) # 转换网络 transformer = TransformerNet() if opt.model_path: transformer.load_state_dict( t.load(opt.model_path, map_location=lambda _s, _: _s)) # 损失网络 Vgg16 置为预测模式 vgg = Vgg16().eval() # 优化器(需要训练 风格转化网络的参数) optimizer = t.optim.Adam(transformer.parameters(), opt.lr) # 获取风格图片的数据 形状 1*c*h*w, 分布 -2~2(使用预设) style = utils.get_style_data(opt.style_path) # 可视化风格图:-2 到2 转化为0-1 vis.img('style', (style[0] * 0.225 + 0.45).clamp(min=0, max=1)) if opt.use_gpu: transformer.cuda() style = style.cuda() vgg.cuda() # 风格图片的gram矩阵 style_v = Variable(style, volatile=True) # 得到vgg中间四层的结果(用以跟输入图片的输出四层比较,计算损失) features_style = vgg(style_v) # gram_matrix:输入 b,c,h,w 输出 b,c,c 计算gram矩阵(四层的gram矩阵) gram_style = [Variable(utils.gram_matrix(y.data)) for y in features_style] # 损失统计 仪表盘 用以可视化(每个epoch中的所有batch平均损失) # 风格损失 style_meter = tnt.meter.AverageValueMeter() # 内容损失 content_meter = tnt.meter.AverageValueMeter() for epoch in range(opt.epoches): # 仪表盘清零 content_meter.reset() style_meter.reset() for ii, (x, _) in tqdm.tqdm(enumerate(dataloader)): # 训练 optimizer.zero_grad() if opt.use_gpu: x = x.cuda() # x为输入的真实图像 x = Variable(x) # 风格转换后的预测图像为y y = transformer(x) # 输入: b, ch, h, w 0~255 # 输出: b, ch, h, w - 2~2 # 将x,y范围从0-255转化为-2-2 y = utils.normalize_batch(y) x = utils.normalize_batch(x) # 返回 四个中间层的特征输出 features_y = vgg(y) features_x = vgg(x) # content loss内容损失 只计算relu2_2之间的损失 预测图片与原图在relu2_2中间层比较,计算损失 # content_weight内容的权重 mse_loss均方误差损失函数 content_loss = opt.content_weight * F.mse_loss( features_y.relu2_2, features_x.relu2_2) # style loss style_loss = 0. # 风格损失取四层的均方误差损失总和 # features_y:预测图像的四层输出内容 gram_style:风格图像的四层输出的gram_matrix # zip将可迭代的对象作为参数,将对象中对应的元素打包成一个个元组,然后返回由这些元组组成的列表 for ft_y, gm_s in zip(features_y, gram_style): # 计算预测图像的四层输出内容的gram_matrix gram_y = utils.gram_matrix(ft_y) style_loss += F.mse_loss(gram_y, gm_s.expand_as(gram_y)) style_loss *= opt.style_weight # 总损失=风格损失+内容损失 total_loss = content_loss + style_loss # 反向传播 total_loss.backward() # 更新参数 optimizer.step() # 损失平滑 将损失加入仪表盘,以便可视化损失过程 content_meter.add(content_loss.data[0]) style_meter.add(style_loss.data[0]) # 每plot_every次前向传播后可视化 if (ii + 1) % opt.plot_every == 0: if os.path.exists(opt.debug_file): ipdb.set_trace() # 可视化 vis.plot('content_loss', content_meter.value()[0]) vis.plot('style_loss', style_meter.value()[0]) # 因为x和y经过标准化处理(utils.normalize_batch),所以需要将它们还原 #x,y为[-2,2]还原回[0,1] vis.img('output', (y.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) vis.img('input', (x.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) # 每次epoch完毕后保存visdom和模型 vis.save([opt.env]) t.save(transformer.state_dict(), 'checkpoints/%s_style.pth' % epoch)
def train(**kwargs): opt = Config() for k_, v_ in kwargs.items(): setattr(opt, k_, v_) device = t.device('cuda') if opt.use_gpu else t.device('cpu') vis = utils.Visualizer(opt.env) # 数据加载 transfroms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Lambda(lambda x: x * 255) ]) dataset = tv.datasets.ImageFolder(opt.data_root, transfroms) dataloader = data.DataLoader(dataset, opt.batch_size) # 转换网络 transformer = TransformerNet() if opt.model_path: transformer.load_state_dict( t.load(opt.model_path, map_location=lambda _s, _: _s)) transformer.to(device) # 损失网络 Vgg16 vgg = Vgg16().eval() vgg.to(device) for param in vgg.parameters(): param.requires_grad = False # 优化器 optimizer = t.optim.Adam(transformer.parameters(), opt.lr) # 获取风格图片的数据 style = utils.get_style_data(opt.style_path) vis.img('style', (style.data[0] * 0.225 + 0.45).clamp(min=0, max=1)) style = style.to(device) # 风格图片的gram矩阵 with t.no_grad(): features_style = vgg(style) gram_style = [utils.gram_matrix(y) for y in features_style] # 损失统计 style_meter = tnt.meter.AverageValueMeter() content_meter = tnt.meter.AverageValueMeter() for epoch in range(opt.epoches): content_meter.reset() style_meter.reset() for ii, (x, _) in tqdm.tqdm(enumerate(dataloader)): # 训练 optimizer.zero_grad() x = x.to(device) y = transformer(x) y = utils.normalize_batch(y) x = utils.normalize_batch(x) features_y = vgg(y) features_x = vgg(x) # content loss content_loss = opt.content_weight * F.mse_loss( features_y.relu2_2, features_x.relu2_2) # style loss style_loss = 0. for ft_y, gm_s in zip(features_y, gram_style): gram_y = utils.gram_matrix(ft_y) style_loss += F.mse_loss(gram_y, gm_s.expand_as(gram_y)) style_loss *= opt.style_weight total_loss = content_loss + style_loss total_loss.backward() optimizer.step() # 损失平滑 content_meter.add(content_loss.item()) style_meter.add(style_loss.item()) if (ii + 1) % opt.plot_every == 0: if os.path.exists(opt.debug_file): ipdb.set_trace() # 可视化 vis.plot('content_loss', content_meter.value()[0]) vis.plot('style_loss', style_meter.value()[0]) # 因为x和y经过标准化处理(utils.normalize_batch),所以需要将它们还原 vis.img('output', (y.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) vis.img('input', (x.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) # 保存visdom和模型 vis.save([opt.env]) t.save(transformer.state_dict(), 'checkpoints/%s_style.pth' % epoch)
folder_num = 8 valid_folder_num = 5 num_epoch = 100 use_cuda = torch.cuda.is_available() path = os.path.expanduser('~/codedata/ice/') print('loading data.....') images_all, labels_all, inc_angle_all = read_clean(path, 'train_clean_size.json') train_set_folders = train_cross(images_all, labels_all, inc_angle_all, folder_num) best_test_loss_stl = [np.inf] * folder_num best_train_loss_stl = [np.inf] * folder_num vis = utils.Visualizer(env='lxg') for folder in range(folder_num): if folder is valid_folder_num: break train_data, train_label, train_inc, test_data, test_label, test_inc = \ train_set_folders.getset(folder) train_dataset = DataSet(train_data, train_label, train_inc, train=True) train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=5) test_dataset = DataSet(test_data, test_label, test_inc, train=False) test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size,
def start_visdom(self): self.vis = utils.Visualizer(env='Adversarial AutoEncoder Training', port=8888)
def start_visdom(self): self.vis = utils.Visualizer(port=8888)
''' 返回为数据集中所有图片的个数 :return: ''' return len(self.imgs) if __name__ == '__main__': picFulPath( '/home/mlxuan/project/DeepLearning/data/benchmark/benchmark_RELEASE/dataset/train.txt', '/home/mlxuan/project/DeepLearning/data/benchmark/benchmark_RELEASE/dataset/img/', '/home/mlxuan/project/DeepLearning/data/benchmark/benchmark_RELEASE/dataset/cls/' ) train_dataset = SBDClassSeg('./ImagAndLal.txt') from torch.utils.data import DataLoader trainloader = DataLoader(train_dataset, batch_size=2, shuffle=False, drop_last=True) V = utils.Visualizer() #按batch_size遍历trainloader for i, (data, label) in enumerate( trainloader ): # 如果batch_size不是1,那么每次输入的batch_size个Tensor就需要有相同的size,否则就不能遍历。还可以用for data,label in trainloader:遍历 data, label = trainloader.dataset.untransforms(data, label) V.plot_many('imgs' + str(i), data) V.plot_img('imgs' + str(i), label) V.close()
def train(**kwargs): opt = Config() for _k, _v in kwargs.items(): setattr(opt, _k, _v) device = t.device("cuda" if t.cuda.is_available() else "cpu") vis = utils.Visualizer(opt.env) # 数据加载 transforms = tv.transforms.Compose([ tv.transforms.Resize(opt.image_size), tv.transforms.CenterCrop(opt.image_size), tv.transforms.ToTensor(), tv.transforms.Lambda(lambda x: x * 255) ]) dataset = tv.datasets.ImageFolder(opt.data_root, transforms) dataloader = data.DataLoader(dataset, opt.batch_size) # 风格转换网络 transformer = TransformerNet() if opt.model_path: transformer.load_state_dict( t.load(opt.model_path, map_location=t.device('cpu'))) transformer.to(device) # 损失网络 Vgg16 vgg = Vgg16().eval() vgg.to(device) for param in vgg.parameters(): param.requires_grad = False # 优化器 optimizer = t.optim.Adam(transformer.parameters(), opt.lr) # 获取风格图片的数据 style = utils.get_style_data(opt.style_path) vis.img('style', (style.data[0] * 0.225 + 0.45).clamp(min=0, max=1)) style = style.to(device) # 风格图片的gramj矩阵 with t.no_grad(): features_style = vgg(style) gram_style = [utils.gram_matrix(y) for y in features_style] # 损失统计 style_loss_avg = 0 content_loss_avg = 0 for epoch in range(opt.epoches): for ii, (x, _) in tqdm(enumerate(dataloader)): # 训练 optimizer.zero_grad() x = x.to(device) y = transformer(x) # print(y.size()) y = utils.normalize_batch(y) x = utils.normalize_batch(x) features_x = vgg(x) features_y = vgg(y) # content loss content_loss = opt.content_weight * F.mse_loss( features_y.relu3_3, features_x.relu3_3) # style loss style_loss = 0 for ft_y, gm_s in zip(features_y, gram_style): with t.no_grad(): gram_y = utils.gram_matrix(ft_y) style_loss += F.mse_loss(gram_y, gm_s.expand_as(gram_y)) style_loss *= opt.style_weight total_loss = content_loss + style_loss total_loss.backward() optimizer.step() content_loss_avg += content_loss.item() style_loss_avg += style_loss.item() if (ii + 1) % opt.plot_every == 0: vis.plot('content_loss', content_loss_avg / opt.plot_every) vis.plot('style_loss', style_loss_avg / opt.plot_every) content_loss_avg = 0 style_loss_avg = 0 vis.img('output', (y.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) vis.img('input', (x.data.cpu()[0] * 0.225 + 0.45).clamp(min=0, max=1)) if (ii + 1) % opt.save_every == 0: vis.save([opt.env]) t.save(transformer.state_dict(), 'checkpoints/%s_style.pth' % (ii + 1))
if use_gpu: net = net.cuda() print(net) train_data = YoloDataloader(img_root = train_img_root,label_path = train_label_path,train = True,transforms = trans) train_dataloader = torch.utils.data.DataLoader(train_data,batch_size= batch_size,shuffle=True,num_workers=4) val_data = YoloDataloader(img_root = val_img_root,label_path = val_label_path,train = False,transforms = trans) val_dataloader = torch.utils.data.DataLoader(val_data,batch_size= batch_size,shuffle=False,num_workers=4) ## define yolov1 loss criterion = Lossv1(Ceils,Box,coor_l,noor_l) ## only optimize full connect weights optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad,net.parameters()), lr = learning_rate,momentum=momentum,weight_decay=decay) vis = utils.Visualizer(env="handsome") ## visdom def train(epoch): net.train() total_loss = 0. for index,(imgs,labels) in enumerate(train_dataloader): if use_gpu: imgs = imgs.cuda() labels = labels.cuda() imgs = Variable(imgs) labels = Variable(labels) optimizer.zero_grad() out = net(imgs) out = out.view(batch_size,Ceils,Ceils,30) loss = criterion(out,labels)