파트 9: Rust로 Simba Network 생성
Source: Dev.to
이미지 → CSV 인코더 / 디코더
나는 사자 새끼의 흑백 그림을 찾아서 다음 스크립트를 사용해 픽셀을 CSV 숫자로 (그리고 다시) 변환했다:
Image to CSV Encoder / Decoder
CSV가 생성된 후, 평소 사용하던 Python 훈련 스크립트로 신경망에 입력했다.
머신이 그림을 그리도록 돕기
스크립트가 여러 지점에서 문제를 일으켰기 때문에, 심바를 그리는 방법을 네트워크가 학습하도록 몇 가지 수정을 추가했습니다.
대형 이미지 문제
- 원본 크기: 474 × 474 px – 훈련 속도가 견딜 수 없을 정도로 느려짐.
- 조정된 크기: 200 × 200 px – 훈련 시간을 합리적으로 유지함.
머신 충돌 및 재시작
가끔 과열로 인해 머신이 충돌하고, 훈련이 매번 처음부터 다시 시작되어 시간과 자원이 크게 낭비되었습니다.
Solution: 마지막 저장된 모델에서 훈련을 재개할 수 있도록 save / load checkpoint 메커니즘을 추가했습니다.
오류‑진동 문제
학습률이 아무리 작아도 손실이 진동 루프에 갇혔습니다. 매우 작은 학습률(예: 0.000001)을 시도했지만, 학습 속도가 비현실적으로 느려졌습니다.
대신 코사인 애닐링을 구현하여 학습률을 점진적으로 감소시켰습니다:
# Cosine annealing learning‑rate schedule
decay_factor = 0.5 * (1 + math.cos(math.pi * i / (epochs + epoch_offset)))
current_lr = lr_min + (lr_max - lr_min) * decay_factor이렇게 하면 부드럽고 안정적인 학습이 이루어졌습니다.
Source:
결과: 수학을 통한 예술
문제를 수정하고 ≈ 1.2 M 회(내 컴퓨터에서는 약 4 시간) 동안 네트워크를 실행한 후, 생성된 출력은 입력 이미지와 매우 가깝게 나왔습니다.
입력은 매우 복잡하고 고차원적인 함수이며, 단순한 XOR 게이트나 로지스틱 회귀 데이터셋을 훨씬 넘어섭니다. 이 실험은 보편 근사 정리를 보여줍니다: 신경망은 (거의) 모든 연속 함수를 표현할 수 있습니다.
원본 이미지
초기 정적(무작위 가중치)
최종 재구성 (200 × 200)
높은 해상도에서의 재구성
학습된 가중치를 사용하여 네트워크를 다양한 캔버스 크기(50 × 50, 512 × 512, 1024 × 1024 등)에서 테스트했습니다. 이미지가 일관되게 재현되어, 네트워크가 개별 픽셀을 외우는 것이 아니라 근본적인 함수를 학습했음을 보여줍니다.
1024 × 1024 재구성
네트워크가 선들의 수학적 개념을 학습했기 때문에, 1024 × 1024 버전은 단순히 확대한 것처럼 픽셀화되지 않고, 마치 네트워크가 더 큰 캔버스에 걸작을 다시 그리는 듯한 모습을 보여줍니다.
검증 및 다음 단계
나는 매우 복잡하고 비효율적이며 비용이 많이 드는 이미지 스케일러를 만들었다. 결과는 만족스럽지만 완벽하지는 않다—핵심은 증명했지만, 더 높은 충실도를 원한다.
Reddit에 결과를 올린 뒤, 한 사용자가 SIREN(Sinusoidal Representation Networks)을 시도해 보라고 제안했다.
- SIREN은 ReLU/Tanh 대신 사인 활성화 함수를 사용한다.
- 암시적 신경 표현에 뛰어나며, 이는 내가 시도하고 있던 기술과 매우 유사하다.
나는 이제 Python으로 SIREN을 구현했으며, 더 나은 재구성을 기대한다. 다음 게시물을 기대해 주세요!
러스트 컴백
Python 프로토타입이 좋은 결과를 보여주면서, 러스트로 문제를 해결하려는 열정을 다시 불러일으켰습니다. 짧은 휴식 후, 저는 또 일주일을 투자해 러스트 구현을 준비했습니다:
Tensor트레이트를 생성하고 모든 정의된 메서드를 그 안으로 이동했습니다.Tensor트레이트를 구현한CpuTensor구조체를 구현했습니다.Tensor트레이트를 구현한GpuTensor구조체를 구현했습니다.
초기 충격
GPU 텐서는 90 + 초가 걸려 같은 네트워크를 Python에서는 8 초에 실행했습니다. Nsight Systems(nsys)로 프로파일링한 결과, 실행 시간 대부분이 CUDA 커널이 아니라 메모리 할당·해제에 소비되고 있었습니다.
CuPy는 맞춤형 메모리‑풀 구현을 사용해 이 오버헤드를 피합니다. Rust에서도 비슷한 해결책이 필요했습니다. cust 크레이트는 직접적인 풀 API를 제공하지 않지만, CUDA 드라이버가 제공하는 풀 기능을 cust::sys를 통해 재내보내고 있습니다. 여러 차례의 시행착오 끝에 작동하는 메모리‑풀 헬퍼를 통합했습니다.
메모리‑풀 헬퍼 (Rust)
pub fn get_mem_pool() -> CudaMemoryPool {
// Get the first CUDA device
let device = Device::get_device(0).unwrap();
// Create a memory pool for the device
let mut pool = std::ptr::null_mut();
let pool_props = CUmemPoolProps {
allocType: cust::sys::CUmemAllocationType::CU_MEM_ALLOCATION_TYPE_PINNED,
handleTypes: cust::sys::CUmemAllocationHandleType::CU_MEM_HANDLE_TYPE_NONE,
location: cust::sys::CUmemLocation {
type_: cust::sys::CUmemLocationType_enum::CU_MEM_LOCATION_TYPE_DEVICE,
id: 0,
},
win32SecurityAttributes: std::ptr::null_mut(),
reserved: [0u8; 64],
};
// Create the pool (unsafe because it calls the CUDA driver API)
unsafe { cuMemPoolCreate(&mut pool, &pool_props) };
// Reserve a large chunk of memory once, then return it to the pool.
let reserve_size: usize = 2048 * 1024 * 1024; // 2 GiB
let mut reserve_ptr: CUdeviceptr = 0;
unsafe {
// Allocate from the pool (synchronous the first time)
cuMemAllocFromPoolAsync(
&mut reserve_ptr,
reserve_size,
pool,
std::ptr::null_mut(),
);
cuStreamSynchronize(std::ptr::null_mut());
// Immediately free it back to the pool for reuse
cuMemFreeAsync(reserve_ptr, std::ptr::null_mut());
cuStreamSynchronize(std::ptr::null_mut());
}
println!("Memory pool created for device {}", device.name().unwrap());
CudaMemoryPool {
pool: Arc::new(Mutex::new(UnsafeCudaMemPoolHandle(pool))),
}
}이 헬퍼를 간단한 main 프로그램에서 실행하면 수천 개의 메모리 블록이 밀리초 단위로 할당·해제되는 것을 확인할 수 있습니다. 한 번만 수행되는 풀 생성 비용은 이후 할당에서 얻는 속도 향상에 비해 무시할 수준입니다.
GpuTensor 통합
헬퍼를 추가한 뒤, Arc<…>를 사용해 풀을 유지하고 모든 디바이스 할당을 그 풀을 통해 감쌌습니다. 처음에는 성능 문제가 계속되었기 때문에 더 조사했습니다.
복사, 붙여넣기… 그리고… 컴파일러 오류
풀은 프로그램 전체 수명 동안 살아남도록 CUDA 컨텍스트와 유사하게 전역에 저장해야 했습니다. 또한 CUDA 래퍼 자체가 여전히 기본 할당자를 사용하고 있음을 발견했으며, 따라서 해당 호출들을 풀 기반 할당으로 교체해야 했습니다.
원시 포인터를 다루는 것이 필요해졌습니다. 아래는 호스트 메모리와 디바이스 메모리를 연결하는 커스텀 디바이스 버퍼의 간소화된 예시입니다:
impl Drop for CustomDeviceBuffer {
fn drop(&mut self) {
let pool = match &GPU_CONTEXT.get() {
// …
};
// …
}
}(구현 세부 사항은 간략히 생략했습니다.)
회고
- SIREN vs. Sigmoid – 사인 활성화 함수로 전환하면서 재구성이 더 선명해졌지만, 학습 곡선이 급격했다.
- Python → Rust – 파이썬 프로토타입이 자신감을 주었고, 러스트는 특히 메모리 관리와 관련된 성능 함정을 드러내어 CUDA 내부 구조에 대한 깊은 탐구를 촉발했다.
- Memory Pools –
cust::sys를 통한 CUDA 메모리‑풀 API 사용으로 할당 오버헤드가 크게 감소했다. 풀을 (Arc<…>) 살아 있게 유지하고 모든 디바이스 할당을 이를 통해 라우팅하는 것이 핵심이었다. - Raw Pointers – 신중한
Drop구현과 적절한 동기화로 안전하게 다룰 수 있다.
예시: 전역 컨텍스트에서 CUDA 풀 가져오기 (Rust)
// 풀을 가져옵니다 (컨텍스트나 풀이 없으면 패닉 발생)
let pool = GPU_CONTEXT
.get()
.expect("No GPU context set")
.pool
.as_ref()
.expect("CUDA not initialized or GPU pool not set");예시: 디바이스 포인터 해제 (Rust)
// `pool` is the `Arc<MemoryPool>` obtained above
let _ = pool.free(self.as_device_ptr().as_raw());예시: 풀을 통한 커스텀 디바이스 버퍼 할당 (Rust)
use cust::{
memory::{DeviceBuffer, DevicePointer},
prelude::*,
};
use std::{
mem::size_of,
sync::Arc,
};
/// Wrapper around a `DeviceBuffer` that is always allocated from the CUDA memory pool.
pub struct CustomDeviceBuffer<T> {
pub device_buffer: DeviceBuffer<T>,
}
impl<T> CustomDeviceBuffer<T> {
/// Allocate a buffer of `size` elements using the global CUDA memory pool.
pub fn get_device_buffer(size: usize) -> Self {
// 1️⃣ Grab the pool (panic if it isn’t available)
let pool: Arc<cust::memory::MemoryPool> = GPU_CONTEXT
.get()
.expect("No GPU context set")
.pool
.as_ref()
.expect("CUDA not initialized or GPU pool not set")
.clone();
// 2️⃣ Compute the byte size, checking for overflow
let byte_size = size
.checked_mul(size_of::<T>())
.expect("Requested allocation size overflowed");
if byte_size == 0 {
panic!("Attempted to allocate a zero‑size buffer");
}
// 3️⃣ Allocate raw memory from the pool
let raw_ptr = pool
.allocate(byte_size)
.expect("CUDA memory‑pool allocation failed");
// 4️⃣ Turn the raw pointer into a `DevicePointer<T>`
let dev_ptr = unsafe { DevicePointer::from_raw(raw_ptr as *mut T) };
// 5️⃣ Build a `DeviceBuffer<T>` from the raw parts (unsafe but safe here)
let device_buffer = unsafe { DeviceBuffer::from_raw_parts(dev_ptr, size) };
// 6️⃣ Return the wrapped buffer
Self { device_buffer }
}
}핵심 포인트
- 안전성 –
DevicePointer로의 원시 포인터 변환과DeviceBuffer로의 원시 파트 변환에만unsafe블록이 사용됩니다. 모든 전제 조건(0이 아닌 크기, 올바른 정렬, 성공적인 할당)은 사전에 검증됩니다. - 메모리 풀 사용 – 할당은
pool.allocate를 통해 수행되어 버퍼가 풀의 재사용 전략에 참여하도록 보장합니다. - 오류 처리 – 간결함을 위해
expect를 사용했으며, 실제 프로덕션 코드에서는 패닉 대신Result<_, cust::error::CudaError>를 반환하도록 구현해야 합니다.
원시 포인터 크기 문제
발생한 일
- 배열의 길이만을 기준으로 메모리를 할당했습니다.
- 요소 타입의 바이트 크기를 곱하는 것을 잊었습니다.
길이가 1인 f32 배열의 경우, 필요한 4 바이트 대신 1 바이트만 할당했습니다.
해결 방법
// 잘못된 할당 (길이만)
let size = length; // → f32[1]에 대해 1 바이트
// 올바른 할당 (길이 × 요소 크기)
let size = length * std::mem::size_of::<f32>(); // → f32[1]에 대해 4 바이트할당을 수정한 후, 코드는 성공적으로 컴파일되고 실행되었지만 실행 시간은 개선되지 않았습니다.
숨겨진 버그
Tensor::ones는 매 epoch마다 1들의 Vec를 다시 만들고 장치로 복사하고 있었다. 나는 이를 메모리를 직접 채우는 GPU 커널로 교체했다:
extern "C" __global__ void fill_value(float *out, int n, float value)
{
int idx = blockIdx.x * blockDim.x + threadIdx.x;
if (idx < n) {
out[idx] = value;
}
}cuBLAS 래퍼 (Rust)
Result {
let m = self.shape[0] as i32;
let k = self.shape[1] as i32;
let n = rhs.shape[1] as i32;
let total_elements = (m * n) as usize;
let result = get_device_buffer(total_elements);
let alpha = T::one();
let beta = T::zero();
unsafe {
cublasSgemm_v2(
Self::_get_cublas_handle(),
cublasOperation_t::CUBLAS_OP_N,
cublasOperation_t::CUBLAS_OP_N,
n,
m,
k,
&alpha.f32(),
rhs.device_buffer.as_device_ptr().as_raw() as *const f32,
n,
self.device_buffer.as_device_ptr().as_raw() as *const f32,
k,
&beta.f32(),
result.as_device_ptr().as_raw() as *mut f32,
n,
);
}
let result_shape = vec![self.shape[0], rhs.shape[1]];
Ok(Self::_with_device_buffer(result_shape, result))
}cuBLAS 버전은 행렬 크기가 보통이었기 때문에 속도 향상이 없었다; cuBLAS는 매우 큰 행렬에서 빛을 발한다.
결론
여정은 길고 좌절이 많았지만, 각 장애물은 신경망 설계와 Rust에서의 저수준 GPU 프로그래밍에 대한 귀중한 교훈을 주었습니다. 이제 저는 CUDA 메모리 풀, 맞춤형 디바이스 버퍼, 그리고 (선택적으로) 고품질 이미지 재구성을 위한 SIREN 네트워크를 활용하는 기능적인 Rust Tensor 라이브러리를 가지고 있습니다.
Note: XOR 테스트 데이터셋은 여기서는 문제가 되지 않으며; 솔루션은 작동하고, 훨씬 큰 이미지 재구성 작업을 수행할 때 확실히 도움이 될 것입니다.