比赛解决方案1-微调Faster-RCNN about global-wheat-detection HOT 5 OPEN

for epo in range(epoch):
        loss_hist.reset()
        val_loss_hist.reset()
        model.train()
        for images, targets, image_ids in train_data_loader:
            
            images = list(image.to(device) for image in images)
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
            # 这里应该是model中封装了求loss，官方教程中写的是 model(images, targets):Returns losses and detections
            loss_dict = model(images, targets)  # Returns losses and detections
            # 对一个batch的loss求和
            losses = sum(loss for loss in loss_dict.values())
            # item方法是取出数值
            loss_value = losses.item()
            # 使用封装的类记录loss
            loss_hist.send(loss_value)

            optimizer.zero_grad()
            losses.backward()
            optimizer.step()

            if itr % 50 == 0:
                print(f"Iteration #{itr} loss: {loss_value}")
            itr += 1

        # 验证
        model.eval()
        with torch.no_grad():
            for val_images, val_targets, val_image_ids in valid_data_loader:
                images = list(image.to(device) for image in val_images)
                targets = [{k: v.to(device) for k, v in t.items()} for t in val_targets]

                val_loss_dict = model(images, targets)  # Returns losses and detections
                val_losses = sum(loss for loss in val_loss_dict.values())
                loss_value = losses.item()
                loss_hist.send(loss_value)

        # 判断是否为更优的模型，以loss为标准
        if val_loss_hist.value<least_loss:
            least_loss = val_loss_hist.value
            lval=int(least_loss*1000)/1000
            torch.save(model.state_dict(), f'fasterrcnn_custom_test_ep{epo}_loss{lval}.pth')
            
        else:
            if lr_scheduler is not None:
                lr_scheduler.step()
        print(f"Epoch #{epo} train_loss: {loss_hist.value} val_loss: {val_loss_hist.value}")