unet pytorch

这篇具有很好参考价值的文章主要介绍了unet pytorch。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

1.单机多卡版本:代码中的DistributedDataParallel (DDP) 部分对应单机多卡的分布式训练方式

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader
from torchvision.transforms import RandomHorizontalFlip, RandomVerticalFlip, RandomRotation, RandomResizedCrop, ToTensor
from torch.nn.parallel import DistributedDataParallel as DDP

# 定义ResNet块
class ResNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ResNetBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.relu(out)
        out = self.conv2(out)
        out += residual
        out = self.relu(out)
        return out

# 定义UNet模型
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
        self.block1 = ResNetBlock(64, 64)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.block2 = ResNetBlock(128, 128)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.block3 = ResNetBlock(256, 256)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1
                self.block4 = ResNetBlock(512, 512)
        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.conv5 = nn.Conv2d(128, out_channels, kernel_size=1)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.block1(x1)
        x2 = self.pool1(x1)
        x2 = self.conv2(x2)
        x2 = self.block2(x2)
        x3 = self.pool2(x2)
        x3 = self.conv3(x3)
        x3 = self.block3(x3)
        x4 = self.pool3(x3)
        x4 = self.conv4(x4)
        x4 = self.block4(x4)
        x = self.upconv3(x4)
        x = torch.cat((x, x3), dim=1)
        x = self.conv5(x)
        x = self.upconv2(x)
        x = torch.cat((x, x2), dim=1)
        x = self.upconv1(x)
        x = torch.cat((x, x1), dim=1)
        x = self.conv5(x)
        return x


# 定义数据集类
class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data = # Load data from data_dir
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        image, mask = self.data[index]
        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)
        return image, mask

# 设置训练参数
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_epochs = 10
batch_size = 4

# 创建UNet模型和优化器
model = UNet(in_channels=3, num_classes=2).to(device)
model = DDP(model)

optimizer = optim.Adam(model.parameters(), lr=0.001)

# 定义数据增强方法
transform = transforms.Compose([
    RandomHorizontalFlip(),
    RandomVerticalFlip(),
    RandomRotation(15),
    RandomResizedCrop(256, scale=(0.8, 1.0)),
    ToTensor(),
])

# 加载数据集并进行数据增强
dataset = CustomDataset(data_dir="path_to_dataset", transform=transform)
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

# 训练循环
for epoch in range(num_epochs):
    model.train()
    total_loss = 0.0
    for images, masks in dataloader:
        images = images.to(device)
        masks = masks.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = nn.CrossEntropyLoss()(outputs, masks)
        loss.backward()
        optimizer.step()

        total_loss += loss.item()

    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {total_loss/len(dataloader)}")

2.多机多卡版本:使用torch.utils.data.distributed.DistributedSampler和torch.distributed.init_process_group来实现多机多卡的分布式训练,确保在每个进程中都有不同的数据划分和完整的通信。文章来源地址https://www.toymoban.com/news/detail-657116.html

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torch.nn.parallel import DistributedDataParallel
from torchvision.transforms import transforms
from torchvision.datasets import YourDataset
from torch.utils.data.distributed import DistributedSampler
import torch.distributed as dist

# 定义ResNet块
class ResNetBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(ResNetBlock, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.relu(out)
        out = self.conv2(out)
        out += residual
        out = self.relu(out)
        return out

# 定义UNet模型
class UNet(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(UNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)
        self.block1 = ResNetBlock(64, 64)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.block2 = ResNetBlock(128, 128)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.block3 = ResNetBlock(256, 256)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv4 = nn.Conv2d(256, 512, kernel_size=3, padding=1
                self.block4 = ResNetBlock(512, 512)
        self.upconv3 = nn.ConvTranspose2d(512, 256, kernel_size=2, stride=2)
        self.upconv2 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.upconv1 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        self.conv5 = nn.Conv2d(128, out_channels, kernel_size=1)

    def forward(self, x):
        x1 = self.conv1(x)
        x1 = self.block1(x1)
        x2 = self.pool1(x1)
        x2 = self.conv2(x2)
        x2 = self.block2(x2)
        x3 = self.pool2(x2)
        x3 = self.conv3(x3)
        x3 = self.block3(x3)
        x4 = self.pool3(x3)
        x4 = self.conv4(x4)
        x4 = self.block4(x4)
        x = self.upconv3(x4)
        x = torch.cat((x, x3), dim=1)
        x = self.conv5(x)
        x = self.upconv2(x)
        x = torch.cat((x, x2), dim=1)
        x = self.upconv1(x)
        x = torch.cat((x, x1), dim=1)
        x = self.conv5(x)
        return x

# 定义数据集类
class CustomDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data = # Load data from data_dir
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        image, mask = self.data[index]
        if self.transform:
            image = self.transform(image)
            mask = self.transform(mask)
        return image, mask

def main(rank, world_size):
    # 设置分布式训练参数
    torch.cuda.set_device(rank)
    torch.distributed.init_process_group(backend='nccl', init_method='env://', world_size=world_size, rank=rank)

    # 设置训练参数
    num_epochs = 10
    batch_size_per_gpu = 4

    # 创建UNet模型和优化器
    in_channels = 3
   
    model = UNet(in_channels=3, num_classes=2).cuda(rank)
    model = DistributedDataParallel(model, device_ids=[rank])
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 数据增强方法
    transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.RandomRotation(30),
        transforms.RandomResizedCrop(256, scale=(0.8, 1.2)),
        transforms.ToTensor()
    ])

    # 加载训练集和验证集
    train_dataset = CustomDataset(transform=transform)
    train_sampler = DistributedSampler(train_dataset)
    train_loader = DataLoader(train_dataset, batch_size=batch_size_per_gpu, sampler=train_sampler)

    # 训练循环
    for epoch in range(num_epochs):
        model.train()
        total_loss = 0.0

        for images, masks in train_loader:
            images = images.cuda(rank)
            masks = masks.cuda(rank)

            # 执行前向传播和反向传播
            optimizer.zero_grad()
            outputs = model(images)
            loss = F.binary_cross_entropy_with_logits(outputs, masks)
            loss.backward()
            optimizer.step()

            total_loss += loss.item()

        if world_size > 1:
            torch.distributed.all_reduce(total_loss)
        total_loss /= len(train_sampler)
        print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {total_loss:.4f}")

def main_multi_gpu():
    world_size = torch.cuda.device_count()
    if world_size > 1:
        torch.multiprocessing.spawn(main, args=(world_size,), nprocs=world_size, join=True)
    else:
        main(0, 1)

if __name__ == '__main__':
    main_multi_gpu()

到了这里,关于unet pytorch的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 探索人工智能:深度学习、人工智能安全和人工智能编程(文末送书)

    人工智能知识对于当今的互联网技术人来说已经是刚需。但人工智能的概念、流派、技术纷繁复杂,选择哪本书入门最适合呢? 这部被誉为人工智能“百科全书”的《人工智能(第3版)》,可以作为每个技术人进入 AI 世界的第一本书。 购书链接,限时特惠5折 这本书是美国

    2024年02月03日
    浏览(121)
  • 人工智能深度学习

    目录 人工智能 深度学习 机器学习 神经网络 机器学习的范围 模式识别 数据挖掘 统计学习 计算机视觉 语音识别 自然语言处理 机器学习的方法 回归算法 神经网络 SVM(支持向量机) 聚类算法 降维算法 推荐算法 其他 机器学习的分类 机器学习模型的评估 机器学习的应用 机

    2024年02月22日
    浏览(58)
  • 人工智能之深度学习

    第一章 人工智能概述 1.1人工智能的概念和历史 1.2人工智能的发展趋势和挑战 1.3人工智能的伦理和社会问题 第二章 数学基础 1.1线性代数 1.2概率与统计 1.3微积分 第三章 监督学习 1.1无监督学习 1.2半监督学习 1.3增强学习 第四章 深度学习 1.1神经网络的基本原理 1.2深度学习的

    2024年02月09日
    浏览(55)
  • 人工智能、机器学习、深度学习的区别

    人工智能涵盖范围最广,它包含了机器学习;而机器学习是人工智能的重要研究内容,它又包含了深度学习。 人工智能是一门以计算机科学为基础,融合了数学、神经学、心理学、控制学等多个科目的交叉学科。 人工智能是一门致力于使计算机能够模拟、模仿人类智能的学

    2024年02月08日
    浏览(56)
  • 人工智能学习07--pytorch14--ResNet网络/BN/迁移学习详解+pytorch搭建

    亮点:网络结构特别深 (突变点是因为学习率除0.1?) 梯度消失 :假设每一层的误差梯度是一个小于1的数,则在反向传播过程中,每向前传播一层,都要乘以一个小于1的误差梯度。当网络越来越深的时候,相乘的这些小于1的系数越多,就越趋近于0,这样梯度就会越来越小

    2023年04月11日
    浏览(167)
  • 人工智能的深度学习如何入门

    人工智能深度学习近年来成为热门的技术领域,被广泛应用于许多领域,如自然语言处理、图像识别、机器翻译等。学习人工智能深度学习需要具备一定的数学和编程基础,但对于初学者来说,并不需要过于复杂的数学和编程知识。本文将介绍人工智能深度学习的基本概念和

    2024年03月27日
    浏览(63)
  • 深度学习:探索人工智能的前沿

    人工智能(Artificial Intelligence,简称AI)是一门研究如何使计算机能够执行通常需要人类智能的任务的领域。从早期的符号推理到现代的深度学习,人工智能经历了漫长的发展过程。 20世纪50年代,AI的奠基性工作开始,研究者们试图通过符号推理来模拟人类思维过程。然而,

    2024年01月19日
    浏览(75)
  • 一探究竟:人工智能、机器学习、深度学习

    1.1 人工智能是什么?          1956年在美国Dartmounth 大学举办的一场研讨会中提出了人工智能这一概念。人工智能(Artificial Intelligence),简称AI,是计算机科学的一个分支,它企图了解智能的实质,并生产出一种新的能以人类智能相似的方式做出反应的智能机器,该领域的

    2024年02月17日
    浏览(53)
  • 12、人工智能、机器学习、深度学习的关系

    很多年前听一个机器学习的公开课,在QA环节,一个同学问了老师一个问题“ 机器学习和深度学习是什么关系 ”? 老师先没回答,而是反问了在场的同学,结果问了2-3个,没有人可以回答的很到位,我当时也是初学一脸懵,会场准备的小礼品也没有拿到。 后来老师解释“机

    2024年02月05日
    浏览(72)
  • 机器学习、人工智能、深度学习三者的区别

    目录 1、三者的关系 2、能做些什么 3、阶段性目标 机器学习、人工智能(AI)和深度学习之间有密切的关系,它们可以被看作是一种从不同层面理解和实现智能的方法。 人工智能(AI):人工智能是一门研究如何使计算机能够模仿人类智能的学科。它涵盖了各种技术和方法,

    2024年02月14日
    浏览(61)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包