Skip to content

Implementation of `Fully Convolutional Networks for Semantic Segmentation` by Jonathan Long, Evan Shelhamer, Trevor Darrell, UC Berkeley

License

Notifications You must be signed in to change notification settings

xevolesi/pytorch-fcn

Repository files navigation

pytorch-fcn

Implementation of Fully Convolutional Networks for Semantic Segmentation by Jonathan Long, Evan Shelhamer, Trevor Darrell, UC Berkeley

Results

Paper version

I tried to reproduce the results from the journal version of paper, but I did not succeed completely. I think this is because I used vgg16 from torchvision as a backbone and trained with a mini-batches without gradient accumulation (although, in theory, this should be almost the same). Also, because of i trained with mini-batches, the learning rate is slightly increased relative to the one from paper. Here are the results:

Architecture mIOU PyTorch checkpoints
FCN32s 0.62457 link
FCN16s 0.64217 link
FCN8s 0.64389 link

Actually i suppose that you can easily beat my results just by adding more augmentations or EMA model, so do it if you want :) In predefined_configs you will find configs for FCN32-FCN8 to reproduce my results.

Integration with timm package

I've added integration with awesome timm repo so now you can use backbones from it. Please see predefined_configs/config_fcn32_densenet121.yml for example of usage densenet121 backbone from timm. Some results:

Architecture Backbone Additional FCN head mIOU PyTorch checkpoints
FCN32s densenet121 False 0.65869 link
FCN32s densenet121 True 0.66486 link
FCN16s densenet121 True 0.69617 link
FCN8s densenet121 True 0.69778 link
FCN32s tf_efficientnetv2_s_in21ft1k True 0.73326 link
FCN16s tf_efficientnetv2_s_in21ft1k True 0.75485 link
FCN8s tf_efficientnetv2_s_in21ft1k True 0.75771 link

Fixed batch examples

FCN32s

FCN32s predictions for fixed batch

FCN16s

FCN32s predictions for fixed batch

FCN8s

FCN32s predictions for fixed batch

FCN32s-densenet121

FCN32s-densenet121 predictions for fixed batch

FCN16s-densenet121

FCN16s-densenet121 predictions for fixed batch

FCN8s-densenet121

FCN8s-densenet121 predictions for fixed batch

FCN32s-efficientnetv2_s

FCN32s-efficientnetv2_s predictions for fixed batch

FCN16s-efficientnetv2_s

FCN16s-efficientnetv2_s predictions for fixed batch

FCN8s-efficientnetv2_s

FCN8s-efficientnetv2_s predictions for fixed batch

Dataset

PASCAL VOC

For validation:

cd path/to/data
wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
tar -xvf VOCtrainval_11-May-2012.tar

and this image set. Just create seg11valid.txt in path/to/data/VOCdevkit/VOC2012/ImageSets/Segmentation.

For training:

cd path/to/data
wget http://www.eecs.berkeley.edu/Research/Projects/CS/vision/grouping/semantic_contours/benchmark.tgz
mv benchmark.tgz benchmark.tar
tar -xvf benchmark.tar

As i understood this is exact setup as in paper.

Custom dataset

I suggest to translate your dataset to PASCAL VOC format and then just use this repo as is to train on your own dataset, but of course you can write you own dataset implementation and use it.

Installation

You need Python==3.10 and PyTorch==2.0.0 with CUDA==11.8, but i think it's easy to run with other versions of PyTorch. Note that i was training the models with:

  1. CPU: AMD RYZEN 9 7950;
  2. GPU: NVidia RTX 4090 24GB;
  3. RAM: 128 GB.

So my settings may not be acceptable for your configuration.

Core requirements

Core requirements are listed in requirements.txt file. You need them to be able to run training pipeline. The simplest way to install them is to use conda:

conda create -n fcn python=3.10 -y
conda activate fcn
pip install -r requirements.txt

But ofcourse you can use old but gold venv by running the following commands:

python3.10 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt

Optional requirements

I've also provided optional requirements listed in requirements.optional.txt. This file contains:

  1. Weights & biases for logging training process. You don't need them for core functionality, but i strongly recommend you to use this awesome tool. In order to use wandb you need to follow quickstart guide. wandb will need your WANDB_API_KEY environment variable and you can set by creating .env file in the root of this repo (see .env-example);
  2. Packages for exporting to ONNX format.

Development requirements.

These are for development pupropse. They consist of flake8, pytest, and so on. You can read requirements.dev.txt file to get more information about development requirements.

How to train

Directly on local machine

It's quite simple:

  1. Modify config.yml file according to your desires;
  2. Run python train.py.

It shoud works okay.

With docker

First of all you need to install docker and nvidia-container-toolkit to be able to run training inside containers:

  1. To install docker please visit official documentation and don't forget about post-installation steps;
  2. To install nvidia-container-toolkit please follow official installation guide. Note that i performed installation only for nvidia-container-toolkit. I DID NOT install Container Device Interface (CDI).

After installation of docker and nvidia-container-toolkit build the image with docker build -t fcn . command. To run training inside docker container do the following steps:

  1. Modify config.yml accroding to your desires;
  2. Run the following command:
docker run --rm --gpus all --shm-size 100g \

	# This is for training data.
	-v /home/xevolesi/Projects/pytorch-fcn/data/:/fcn/data \

	# This is for logs.
	-v /home/xevolesi/Projects/pytorch-fcn/logs:/fcn/logs \
	
	# Mount configuration file with your own changes.
	--mount type=bind,source="$(pwd)"/config.yml,target=/fcn/config.yml,readonly \
	
	# You need to provide W&B API key to be able to use it.
	# You can ommit it if you don't need it.
	-e WANDB_API_KEY=%your W&B API key% \
	-it fcn

Predict on image

To get predictions on single image please use predict.py script. To run this script just do:

python predict.py \
		--image %path to image% \
		--config %path to config% \
		--weights %path to weights% \
		--image_size %Image height and image width separated by single comma%

Example:

python predict.py \
		--image data/VOCdevkit/VOC2012/JPEGImages/2007_000363.jpg \
		--config config.yml \
		--weights weights/fcn_sim..onnx \
		--image_size 500,500

Export to ONNX

To export the model just do:

python export_to_onnx.py \
		--config %path to config.yml% \
		--torch_weights %path to PyTorch weights% \
		--onnx_path %path tot ONNX file% \
		--image_size %Image height and image width separated by single comma% \
		--do_check_on_validation_set

Example of command to export FCN8:

python export_to_onnx.py \
		--config config.yml \
		--torch_weights logs/2023-04-12\ 08\:04\:42.694413/weights/fcn_8_iou_0.6438923478126526.pt \
		--onnx_path ./fcn.onnx \
		--image_size 500,500 \
		--do_check_on_validation_set

How to develop

Clone this repo and install development dependencies via pip install -r requirements.dev.txt. Makefile consists of the following recipies:

  1. lint - run linter checking;
  2. format - run formatters;
  3. verify_format - run formatters is checking mode;
  4. run_tests - run all tests;
  5. pre_push_test - run formatters and tests. Use make lint to run linter checking and make format to apply formatters.

This repo contains some tests just to be sure that some tricks works correctly, but unfortunatelly it doesn't contain any integration tests (test for the whole modelling pipeline) just because it's really annoying to adapt them to free github runners. Note that running tests takes quite a lot of time.

About

Implementation of `Fully Convolutional Networks for Semantic Segmentation` by Jonathan Long, Evan Shelhamer, Trevor Darrell, UC Berkeley

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published