AI

pytorch 사용 예제

DOAD 2021. 6. 17. 17:51
#라이브러리 불러오기 (torch.nn : 딥러닝 네트워크 구현 및 학습을 간단하게 수행할 수 있도록 다양한 함수 제공)
import torch
import torch.nn as nn
import torch.nn.functional as F

#gpu가 있는 경우 pytorch 연산을 gpu로, 그렇지 않은 경우 pytorch 연산을 cpu로 수행
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#딥러닝 모델
class Model(nn.Module) :
	def __init__(self) :	# 네트워크의 변수 정의
		super(Model, self).__init__()
		self.fc1 = nn.Linear(3, 4)	#nn.Linear(입력수,출력수)
		self.fc2 = nn.Linear(4, 4)
		self.fc3 = nn.Linear(4, 2)
		
	def forward(self, x) :	# 네트워크 구조 결정 및 연산 수행
		x = F.relu(self.fc1(x))	#선형 연산 후 비선형 함수 통과
		x = F.relu(self.fc2(x))
		x = self.fc3(x)	# 선형 연산만 수행
		return x
        
#네트워크 선언
model = Model().to(device)	#Model클래스 호출 후 디바이스 할당하고 이를 model로 선언
	
#Pytorch를 이용한 연산 수행
x = torch.rand(2,3).to(device)	#(2,3)의 랜덤 데이터를 torch로 생성 후 연산장치 할당
				# torch는 pytorch 연산 수행을 위한 데이터 자료형
output = model(x)

#결과 출력
print(output)

#Numpy 변환
output_np = output.cpu().detach().numpy()	# 연산장치를 cpu로 변환
						# detch : gradient 전파가 안되도록 함
                                            # numpy : tensor를 numpy로 변환
print(output_np)

numpy  변환을 위해서는 detach를 꼭 해줘야 함