Treinamento

Baixando a imagem

O primeiro passo é escolher alguma imagem como ponto de partida. No Docker Hub (opens in a new tab) exitem milhares de imagens com os mais diversos pacotes instalados, como a do Ubuntu (opens in a new tab), do Python (opens in a new tab), do banco de dados Postgres (opens in a new tab) ou até mesmo do sistema operacional Alpine (opens in a new tab), muito utilizado para gerar imagens docker de tamanho mais compacto. Todos estes pacotes, possuem, na sua página diversas informações, sobre como usá-lo e sobre suas tags. No nosso caso, iremos utilizar a que os criadores do YOLOv5 oferencem, disponível neste link (opens in a new tab), pois já vêm com todos os pacotes e pré requisitos já instalados.

Para tal, podemos rodar o comando abaixo. Note a ausência de uma tag após o nome da imagem, que fará com que baixamos a imagem com tag latest por padrão. Em muitos casos isso pode gerar problemas quando os autores da imagem a atualizarem, e portanto é recomendado setar alguma outra tag, de forma explcíta, com a versão que você precisa. No nosso caso, porém, isto não será um problema.

docker pull ultralytics/yolov5

Podemos verificar que a imagem foi baixada corretamente com o comando que lista todas as imagens que estão na máquina. O comando segue o mesmo padrão dos comandos Unix, e portanto, rm também é aceito, para remover alguma imagem.

docker image ls

Executando uma imagem

Para rodar uma image, e criar um container, utilizamos o comando run. Porém, no nosso caso, iremos passar também algumas flags que alteram este comando, pois precisamos passar o nosso dataset para dentro do container, e retirar o modelo de inferência de dentro dele, após terminarmos o treinamento.

Lembre que o container criado é um ambiente virtual, que está acima do kernel, e portanto, ele não tem nenhum conhecimento dos arquivos ou programas que temos instalados em nosso computador. Desta forma, é necessário criar um volume (opens in a new tab), que será uma ponte entre nossa máquina, o host, e o container. Todos os arquivos que estiverem dentro da pasta que passarmos serão criados no container, e da mesma forma após criarmos ou editar algum arquivo desta pasta de dentro dele, estas modificações serão repassadas para o host. A 'ponte' criada pelo volume vai, portanto, nas duas direções.

Note que no primeiro volume abaixo, estamos passando a pasta de caminho /home/ubuntu/datasets no host para uma pasta de caiminho /usr/src/datasets dentro do container, com o objetivo de passar nosso dataset para ele. Em contrapartida, fazemos o mapeamento entre /home/ubuntu/podesta/runs, que é uma pasta inicialmente vazia, e /usr/src/app/runs pois este é o local onde o YOLOv5 irá salvar o modelo para inferência.

Passamos ainda as flags -it, que são duas flags, --interactive e --tty, utilizadas para que possamos passar comandos e interagir com o container criado. Além disso, a flag --name é opcional, mas muito importante para que saibamos qual é o nosso container. Caso ela não seja passada, o Docker irá atribuir um nome aleatório para nosso container. O último comando é o nome da imagem.

docker run \
  -it \
  --volume /home/ubuntu/datasets:/usr/src/datasets \
  --volume /home/ubuntu/podesta/runs:/usr/src/app/runs \
  --name train-yolov5 \
  ultralytics/yolov5

Após entrar dentro do container podemos ver que os arquivos não são os mesmos que originalmente estavam em nosso host. Os arquivos são aqueles do repositório do YOLOv5 no Github (opens in a new tab), pois são os arquivos que os criadores da imagem passaram para ela. Podemos ver também, um caminho acima do atual, que nossa pasta datasets também foi corretamente criada. Além disso noso usuário agora é outro, que é um usuário exclusivo do ambiente virtual que foi criado pelo Docker.

Treinando o modelo

Agora que estamos dentro do container, podemos começar a treinar o modelo. Para isso vamos nos guiar pela documentação disponibilizada pelos criadores do modelo, a empresa Ultralytics, sobre como treinar um modelo com custom data (opens in a new tab).

O modelo espera um arquivo com extensão yaml dentro da pasta data, com os metadados do dataset. Vamos então, copiar o nosso yaml, da pasta datasets, para a pasta data que está dentro da pasta app. Podemos então verificar que o arquivo foi de fato movido com o comando ls.

cp ../datasets/road-signs/road-signs.yaml ./data/

Para treinar o modelo vamos utilizar o script train.py, e para isso vamos passar algumas flags para o treinamento. Maiores informações sobre elas podem ser obtidas na documentação oficial do YOLOv5 (opens in a new tab).

python train.py --img 1088 --batch 32 --epochs 50 --data road-signs.yaml --device=0 --weights yolov5s.pt --cache

Teste de Detecção

Para treinar o modelo, vamos utilizar o comando detect.py, passando como --source a pasta onde colocamos nossas imagens de teste, --weights o modelo que acabamos de treinar, e --conf, o valor de confiança necessário para a detecção.

python detect.py  --source /usr/src/datasets/road-signs/test/images --weights runs/train/exp/weights/best.pt --conf 0.5

Podemos então verificar os resultados tanto dentro do container, quanto em nosso host, na pasta runs, já que fizemos esta conexão quando iniciamos nosso container.