def main(class_num, pre_trained, train_data, batch_size, momentum, lr, cate_weight, epoch, weights): device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") model = PSPNet(num_classes=class_num, downsample_factor=16, pretrained=True, aux_branch=False) model = model.to(device) train_loader = Data.DataLoader(dataset=train_data, batch_size=batch_size, shuffle=True) optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum) loss_func = nn.CrossEntropyLoss(weight=torch.from_numpy(np.array(cate_weight)).float()).cuda() model.train() for i in range(epoch): for step, (b_x, b_y) in enumerate(train_loader): b_x = b_x.to(device) b_y = b_y.to(device) b_y = b_y.view(-1, 473, 473) output = model(b_x) loss = loss_func(output, b_y.long()) loss = loss.to(device) optimizer.zero_grad() loss.backward() optimizer.step() if step % 1 == 0: print("Epoch:{0} || Step:{1} || Loss:{2}".format(i, step, format(loss, ".4f"))) torch.save(model.state_dict(), weights + "PSPNet_weights" + ".pth")
model = PSPNet(n_classes=6, pretrained=True).to(device) optimizer = torch.optim.Adam(model.parameters(), lr=0.001) criterion = torch.nn.CrossEntropyLoss() mkdirs = lambda x: os.makedirs(x, exist_ok=True) # model.load_state_dict(torch.load("ckpt_seg/epoch_79_iou0.88.pth")) ckpt_dir = "ckpt_seg" mkdirs(ckpt_dir) epoch = 80 for i in range(0, epoch): # train model.train() epoch_iou = [] epoch_start = time.time() for j, (images, seg_gts, rets) in enumerate(train_loader): images = images.to(device) seg_gts = seg_gts.to(device) optimizer.zero_grad() seg_logit = model(images) loss_seg = criterion(seg_logit, seg_gts.long()) loss = loss_seg loss.backward() optimizer.step() # epoch_loss += loss.item() if j % 10 == 0:
def train(): print(torch.cuda.device_count()) # os.environ["CUDA_VISIBLE_DEVICES"] = '0, 1' net = PSPNet(n_classes=1, sizes=(1, 2, 3, 6), psp_size=512, deep_features_size=256, backend='resnet34', pretrained=False) net = nn.DataParallel(net) net = net.cuda() weight = torch.ones(3) weight[0] = 0 # print(net) optimizer = Adam(net.parameters(), lr=1e-3) # criterion = nn.CrossEntropyLoss() criterion = nn.MSELoss() loaders = prepared_train_data() test_loaders = prepared_test_data() # print(len(loaders)) # print(loaders) for epoch in range(1, 100): print('Training................') epoch_loss = [] iteration = 0 net.train(mode=True) for step, sample_batch in enumerate(loaders): # print("Iter:"+str(iteration)) iteration = iteration + 1 images = sample_batch['image'].cuda() # masks = torch.transpose(sample_batch['mask'], 1, 3) masks = sample_batch['mask'].cuda() # print(images.size()) # print(masks.size()) inputs = Variable(images) targets = Variable(masks) # print(targets.size()) outputs = net(inputs) outputs = torch.clamp(outputs, 0., 255.) # results = outputs.cpu().data.numpy() # # # print(np.shape(results)) # # map = np.squeeze(results, axis=[0, 1]) # # misc.imsave('./test_images/test_image_' + str(iteration) + '.png', map) # print(outputs) # print(outputs) # print(outputs.size()) optimizer.zero_grad() loss = criterion(outputs, targets) loss.backward() optimizer.step() epoch_loss.append(loss.data[0]) # if iteration % 10 == 0: # print("Epoch:{}, iteration:{}, loss:{}".format(epoch, # step, # loss)) # if iteration % 10 == 0: # results = outputs.cpu().data.numpy() # # print(np.shape(results)) # map = np.squeeze(results, axis=1) # # print(np.shape(map)) # # map = np.transpose(map, [0, 2, 3, 1]) # # misc.imsave('./train_image/test_image'+str(iteration)+'.png', map[0, :, :]) # if iteration % 400 == 0: # torch.save(net, 'Models/modol-'+str(epoch)+'-'+str(iteration)+'.pth') # torch.save(net, 'Models/modol-' + str(epoch) + '.pth') print('Testing........................') net.train(mode=False) total_m1 = 0 disc = 44 for iteration, item in enumerate(test_loaders): images = item['image'].cuda() # masks = torch.transpose(sample_batch['mask'], 1, 3) masks = item['mask'].numpy() name = item['name'] masks = np.squeeze(masks, axis=0) test_image = Variable(images).cuda() predict_label = net(test_image) predict_label = torch.clamp(predict_label, 0., 255.) results = predict_label.cpu().data.numpy() map = np.squeeze(results, axis=[0, 1]) gt = np.zeros(shape=np.shape(masks)) gt[masks > 200] = 1 prediction = np.zeros(shape=np.shape(map)) prediction[map > disc] = 1 overlap = gt + prediction # print(overlap.max(), overlap.min()) # print(np.shape(overlap)) image_inter = np.zeros(shape=np.shape(overlap)) image_inter[overlap > 1] = 1 num_inter = sum(sum(image_inter)) # print(np.shape(num_inter)) image_union = np.zeros(shape=np.shape(overlap)) image_union[overlap > 0] = 1 num_union = sum(sum(image_union)) # print(np.shape(num_union)) m_1 = (1 - num_inter / num_union) print('Image name is {}, and m1 is {}'.format(name[0], m_1)) total_m1 = total_m1 + m_1 map[map > disc] = 255 # # misc.imsave('./test_images/test_image_' + str(iteration) + '.png', map) misc.imsave('./test_image/' + name[0], map) print('m1 is {}'.format(total_m1 / 200))
def train(epo_num=50, show_vgg_params=False): vis = visdom.Visdom() device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') vgg_model = VGGNet(requires_grad=True, show_params=show_vgg_params) fcn_model = PSPNet(n_classes=2) fcn_model = fcn_model.to(device) criterion = nn.BCELoss().to(device) optimizer = optim.SGD(fcn_model.parameters(), lr=1e-2, momentum=0.7) all_train_iter_loss = [] all_test_iter_loss = [] # start timing prev_time = datetime.now() for epo in range(epo_num): train_loss = 0 fcn_model.train() for index, (img, lidar, label, color_label) in enumerate(train_dataloader): img = img.to(device) lidar = lidar.to(device) label = label.to(device) optimizer.zero_grad() output, out_cls = fcn_model(img) output = torch.sigmoid( output) # output.shape is torch.Size([4, 2, 160, 160]) loss = criterion(output, label) loss.backward() iter_loss = loss.item() all_train_iter_loss.append(iter_loss) train_loss += iter_loss optimizer.step() output_np = output.cpu().detach().numpy().copy() output_np = np.argmin(output_np, axis=1) bag_msk_np = label.cpu().detach().numpy().copy() bag_msk_np = np.argmin(bag_msk_np, axis=1) if np.mod(index, 15) == 0: print('epoch {}, {}/{},train loss is {}'.format( epo, index, len(train_dataloader), iter_loss)) vis.images(output_np[:, None, :, :], win='train_pred', opts=dict(title='train prediction')) vis.images(bag_msk_np[:, None, :, :], win='train_label', opts=dict(title='label')) vis.line(all_train_iter_loss, win='train_iter_loss', opts=dict(title='train iter loss')) test_loss = 0 fcn_model.eval() for index, (img, lidar, label, color_label) in enumerate(test_dataloader): img = img.to(device) lidar = lidar.to(device) label = label.to(device) with torch.no_grad(): optimizer.zero_grad() output, out_cls = fcn_model(img) output = torch.sigmoid(output) loss = criterion(output, label) iter_loss = loss.item() all_test_iter_loss.append(iter_loss) test_loss += iter_loss output_np = output.cpu().detach().numpy().copy() output_np = np.argmin(output_np, axis=1) bag_msk_np = label.cpu().detach().numpy().copy() bag_msk_np = np.argmin(bag_msk_np, axis=1) if np.mod(index, 15) == 0: print( r'Testing... Open http://localhost:8097/ to see test result.' ) vis.images(output_np[:, None, :, :], win='test_pred', opts=dict(title='test prediction')) vis.images(bag_msk_np[:, None, :, :], win='test_label', opts=dict(title='label')) vis.line(all_test_iter_loss, win='test_iter_loss', opts=dict(title='test iter loss')) cur_time = datetime.now() h, remainder = divmod((cur_time - prev_time).seconds, 3600) m, s = divmod(remainder, 60) time_str = "Time %02d:%02d:%02d" % (h, m, s) prev_time = cur_time print('epoch train loss = %f, epoch test loss = %f, %s' % (train_loss / len(train_dataloader), test_loss / len(test_dataloader), time_str)) if np.mod(epo, 5) == 0: torch.save(fcn_model, 'checkpoints/fcn_model_{}.pt'.format(epo)) print('saveing checkpoints/fcn_model_{}.pt'.format(epo))